Fix/delegate all

bittensor_cli/cli.py

@@ -2670,7 +2670,7 @@ def root_delegate_stake(
             root.delegate_stake(
                 wallet,
                 self.initialize_chain(network, chain),
-                float(amount),
+                amount,
                 delegate_ss58key,
                 prompt,
             )
@@ -2741,7 +2741,7 @@ def root_undelegate_stake(
             root.delegate_unstake(
                 wallet,
                 self.initialize_chain(network, chain),
-                float(amount),
+                amount,
                 delegate_ss58key,
                 prompt,
             )

bittensor_cli/src/commands/root.py

@@ -461,7 +461,7 @@ async def delegate_extrinsic(
     subtensor: SubtensorInterface,
     wallet: Wallet,
     delegate_ss58: str,
-    amount: Balance,
+    amount: Optional[float],
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
     prompt: bool = False,
@@ -472,7 +472,7 @@ async def delegate_extrinsic(
     :param subtensor: The SubtensorInterface used to perform the delegation, initialized.
     :param wallet: Bittensor wallet object.
     :param delegate_ss58: The `ss58` address of the delegate.
-    :param amount: Amount to stake as bittensor balance
+    :param amount: Amount to stake as bittensor balance, None to stake all available TAO.
     :param wait_for_inclusion: If set, waits for the extrinsic to enter a block before returning `True`, or returns
                               `False` if the extrinsic fails to enter the block within the timeout.
     :param wait_for_finalization: If set, waits for the extrinsic to be finalized on the chain before returning `True`,
@@ -484,19 +484,19 @@ async def delegate_extrinsic(
              the response is `True`.
     """
 
-    async def _do_delegation() -> tuple[bool, str]:
+    async def _do_delegation(staking_balance_: Balance) -> tuple[bool, str]:
         """Performs the delegation extrinsic call to the chain."""
         if delegate:
             call = await subtensor.substrate.compose_call(
                 call_module="SubtensorModule",
                 call_function="add_stake",
-                call_params={"hotkey": delegate_ss58, "amount_staked": amount.rao},
+                call_params={"hotkey": delegate_ss58, "amount_staked": staking_balance_.rao},
             )
         else:
             call = await subtensor.substrate.compose_call(
                 call_module="SubtensorModule",
                 call_function="remove_stake",
-                call_params={"hotkey": delegate_ss58, "amount_unstaked": amount.rao},
+                call_params={"hotkey": delegate_ss58, "amount_unstaked": staking_balance_.rao},
             )
         return await subtensor.sign_and_send_extrinsic(
             call, wallet, wait_for_inclusion, wait_for_finalization
@@ -567,16 +567,13 @@ async def get_stake_for_coldkey_and_hotkey(
     # Convert to bittensor.Balance
     if amount is None:
         # Stake it all.
-        staking_balance = Balance.from_tao(my_prev_coldkey_balance.tao)
-    else:
-        staking_balance = amount
-
-    if delegate:
-        # Remove existential balance to keep key alive.
-        if staking_balance > (b1k := Balance.from_rao(1000)):
-            staking_balance = staking_balance - b1k
+        if delegate_string == "delegate":
+            staking_balance = Balance.from_tao(my_prev_coldkey_balance.tao)
         else:
-            staking_balance = staking_balance
+            # Unstake all
+            staking_balance = Balance.from_tao(my_prev_delegated_stake.tao)
+    else:
+        staking_balance = Balance.from_tao(amount)
 
     # Check enough balance to stake.
     if delegate_string == "delegate" and  staking_balance > my_prev_coldkey_balance:
@@ -599,6 +596,16 @@ async def get_stake_for_coldkey_and_hotkey(
         )
         return False
 
+    if delegate:
+        # Grab the existential deposit.
+        existential_deposit = await subtensor.get_existential_deposit()
+
+        # Remove existential balance to keep key alive.
+        if staking_balance > my_prev_coldkey_balance - existential_deposit:
+            staking_balance = my_prev_coldkey_balance - existential_deposit
+        else:
+            staking_balance = staking_balance
+
     # Ask before moving on.
     if prompt:
         if not Confirm.ask(
@@ -615,7 +622,7 @@ async def get_stake_for_coldkey_and_hotkey(
         spinner="aesthetic",
     ) as status:
         print_verbose("Transmitting delegate operation call")
-        staking_response, err_msg = await _do_delegation()
+        staking_response, err_msg = await _do_delegation(staking_balance)
 
     if staking_response is True:  # If we successfully staked.
         # We only wait here if we expect finalization.
@@ -1318,7 +1325,7 @@ async def _do_set_take() -> bool:
 async def delegate_stake(
     wallet: Wallet,
     subtensor: SubtensorInterface,
-    amount: float,
+    amount: Optional[float],
     delegate_ss58key: str,
     prompt: bool,
 ):
@@ -1328,7 +1335,7 @@ async def delegate_stake(
         subtensor,
         wallet,
         delegate_ss58key,
-        Balance.from_tao(amount),
+        amount,
         wait_for_inclusion=True,
         prompt=prompt,
         delegate=True,
@@ -1338,7 +1345,7 @@ async def delegate_stake(
 async def delegate_unstake(
     wallet: Wallet,
     subtensor: SubtensorInterface,
-    amount: float,
+    amount: Optional[float],
     delegate_ss58key: str,
     prompt: bool,
 ):
@@ -1348,7 +1355,7 @@ async def delegate_unstake(
         subtensor,
         wallet,
         delegate_ss58key,
-        Balance.from_tao(amount),
+        amount,
         wait_for_inclusion=True,
         prompt=prompt,
         delegate=False,

tests/e2e_tests/test_root.py

@@ -1,6 +1,7 @@
 import time
 
 from bittensor_cli.src.bittensor.balances import Balance
+from tests.e2e_tests.utils import extract_coldkey_balance
 
 """
 Verify commands:
@@ -26,6 +27,9 @@ def test_root_commands(local_chain, wallet_setup):
         4. Execute list delegates and verify information
         5. Execute set-take command, change the take to 12%, verify
         6. Execute delegate-stake command, stake from Alice to Bob
+        7. Execute undelegate-stake command, unstake from Bob to Alice
+        8. Execute delegate-stake command, stake all from Alice to Bob
+        9. Execute undelegate-stake command, unstake all from Bob to Alice
 
     Raises:
         AssertionError: If any of the checks or verifications fail
@@ -167,8 +171,6 @@ def test_root_commands(local_chain, wallet_setup):
     take_percentage = float(bob_delegate_info[7].strip("%")) / 100
     assert take_percentage == float(new_take)
 
-    delegate_amount = 999999
-
     # Stake to delegate Bob from Alice
     stake_delegate = exec_command_alice(
         command="root",
@@ -185,50 +187,96 @@ def test_root_commands(local_chain, wallet_setup):
             "--network",
             "local",
             "--amount",
-            f"{delegate_amount}",
+            f"10",
             "--no-prompt",
         ],
     )
     assert "✅ Finalized" in stake_delegate.stdout
 
-    # List all delegates of Alice (where she has staked)
-    alice_delegates = exec_command_alice(
+    check_my_delegates(
+        exec_command=exec_command_alice,
+        wallet=wallet_alice,
+        delegate_ss58key=wallet_bob.hotkey.ss58_address,
+        delegate_amount=10
+    )
+
+    check_balance(
+        exec_command=exec_command_alice,
+        wallet=wallet_alice,
+        expected_balance={'free_balance': 999990.0, 'staked_balance': 10.0, 'total_balance': 1000000.0},
+    )
+
+    # TODO: Ask nucleus the rate limit and wait epoch
+    # Sleep 120 seconds for rate limiting when unstaking
+    print("Waiting for interval for 2 minutes")
+    time.sleep(120)
+
+    # Unstake from Bob Delegate
+    undelegate_alice = exec_command_alice(
         command="root",
-        sub_command="my-delegates",
+        sub_command="undelegate-stake",
         extra_args=[
             "--wallet-path",
             wallet_path_alice,
             "--chain",
             "ws://127.0.0.1:9945",
             "--wallet-name",
             wallet_alice.name,
+            "--delegate-ss58key",
+            wallet_bob.hotkey.ss58_address,
             "--network",
             "local",
+            "--amount",
+            f"10",
+            "--no-prompt",
         ],
     )
-    # First row are headers, records start from second row
-    alice_delegates_info = alice_delegates.stdout.splitlines()[5].split()
-
-    # WALLET: Wallet name of Alice
-    assert alice_delegates_info[0] == wallet_alice.name
+    assert "✅ Finalized" in undelegate_alice.stdout
 
-    # SS58: address of the Bob's hotkey (Alice has staked to Bob)
-    assert wallet_bob.hotkey.ss58_address == alice_delegates_info[2]
+    check_balance(
+        exec_command=exec_command_alice,
+        wallet=wallet_alice,
+        expected_balance={'free_balance': 1000000.0, 'staked_balance': 0.0, 'total_balance': 1000000.0},
+    )
 
-    # Delegation: This should be 10 as Alice delegated 10 TAO to Bob
-    delegate_stake = Balance.from_tao(string_tao_to_float(alice_delegates_info[3]))
-    assert delegate_stake == Balance.from_tao(delegate_amount)
+    # TODO: Ask nucleus the rate limit and wait epoch
+    # Sleep 120 seconds for rate limiting when unstaking
+    print("Waiting for interval for 2 minutes")
+    time.sleep(120)
 
-    # TOTAL STAKE(τ): This should be 10 as only Alice has delegated to Bob
-    total_stake = Balance.from_tao(string_tao_to_float(alice_delegates_info[7]))
-    assert total_stake == Balance.from_tao(delegate_amount)
+    # Stake to delegate Bob from Alice
+    stake_delegate = exec_command_alice(
+        command="root",
+        sub_command="delegate-stake",
+        extra_args=[
+            "--wallet-path",
+            wallet_path_alice,
+            "--chain",
+            "ws://127.0.0.1:9945",
+            "--wallet-name",
+            wallet_alice.name,
+            "--delegate-ss58key",
+            wallet_bob.hotkey.ss58_address,
+            "--network",
+            "local",
+            "--all",
+            "--no-prompt",
+        ],
+    )
+    assert "✅ Finalized" in stake_delegate.stdout
 
-    # Total delegated Tao: This is listed at the bottom of the information
-    # Since Alice has only delegated to Bob, total should be 10 TAO
-    total_delegated_tao = Balance.from_tao(
-        string_tao_to_float(alice_delegates.stdout.splitlines()[8].split()[3])
+    # check_my_delegates(
+    #     exec_command=exec_command_alice,
+    #     wallet=wallet_alice,
+    #     delegate_ss58key=wallet_bob.hotkey.ss58_address,
+    #     delegate_amount=999999.9999995
+    # )
+
+    check_balance(
+        exec_command=exec_command_alice,
+        wallet=wallet_alice,
+        expected_balance={'free_balance': 0.0000005, 'staked_balance': 999999.9999995, 'total_balance': 1000000.0},
     )
-    assert total_delegated_tao == Balance.from_tao(delegate_amount)
 
     # TODO: Ask nucleus the rate limit and wait epoch
     # Sleep 120 seconds for rate limiting when unstaking
@@ -250,15 +298,83 @@ def test_root_commands(local_chain, wallet_setup):
             wallet_bob.hotkey.ss58_address,
             "--network",
             "local",
-            "--amount",
-            f"{delegate_amount}",
+            "--all",
             "--no-prompt",
         ],
     )
     assert "✅ Finalized" in undelegate_alice.stdout
 
+    check_balance(
+        exec_command=exec_command_alice,
+        wallet=wallet_alice,
+        expected_balance={'free_balance': 1000000.0, 'staked_balance': 0.0, 'total_balance': 1000000.0},
+    )
+
     print("✅ Passed Root commands")
 
 
+def check_my_delegates(exec_command, wallet, delegate_ss58key, delegate_amount):
+    # List all delegates of Alice (where she has staked)
+    delegates = exec_command(
+        command="root",
+        sub_command="my-delegates",
+        extra_args=[
+            "--wallet-path",
+            wallet.path,
+            "--chain",
+            "ws://127.0.0.1:9945",
+            "--wallet-name",
+            wallet.name,
+            "--network",
+            "local",
+        ],
+    )
+    # First row are headers, records start from second row
+    delegates_info = delegates.stdout.splitlines()[5].split()
+    # WALLET: Wallet name of Alice
+    assert delegates_info[0] == wallet.name
+    # SS58: address of the Bob's hotkey (Alice has staked to Bob)
+    assert delegate_ss58key == delegates_info[2]
+    # Delegation: This should be `delegate_amount` as Alice delegated `delegate_amount` TAO to Bob
+    delegate_stake = Balance.from_tao(string_tao_to_float(delegates_info[3]))
+    assert delegate_stake == Balance.from_tao(delegate_amount)
+    # TOTAL STAKE(τ): This should be `delegate_amount` as only Alice has delegated to Bob
+    total_stake = Balance.from_tao(string_tao_to_float(delegates_info[7]))
+    assert total_stake == Balance.from_tao(delegate_amount)
+    # Total delegated Tao: This is listed at the bottom of the information
+    # Since Alice has only delegated to Bob, total should be `delegate_amount` TAO
+    total_delegated_tao = Balance.from_tao(
+        string_tao_to_float(delegates.stdout.splitlines()[8].split()[3])
+    )
+    assert total_delegated_tao == Balance.from_tao(delegate_amount)
+
+
+def check_balance(exec_command, wallet, expected_balance):
+    # Check balance of Alice after registering to the subnet
+    wallet_balance = exec_command(
+        command="wallet",
+        sub_command="balance",
+        extra_args=[
+            "--wallet-path",
+            wallet.path,
+            "--chain",
+            "ws://127.0.0.1:9945",
+            "--wallet-name",
+            wallet.name,
+            "--network",
+            "local",
+        ],
+    )
+
+    # Extract balance left after creating and registering into the subnet
+    balance = extract_coldkey_balance(
+        wallet_balance.stdout,
+        wallet_name=wallet.name,
+        coldkey_address=wallet.coldkey.ss58_address,
+    )
+
+    assert balance == expected_balance
+
+
 def string_tao_to_float(alice_delegates_info: str) -> float:
     return float(alice_delegates_info.replace(",", "").strip("τ"))