Fixed: delegate-stake and undelegate-stake failed when --all option w…

…as used

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
@@ -569,14 +569,7 @@ async def get_stake_for_coldkey_and_hotkey(
         # 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
-        else:
-            staking_balance = staking_balance
+        staking_balance = Balance.from_tao(amount)
 
     # Check enough balance to stake.
     if delegate_string == "delegate" and  staking_balance > my_prev_coldkey_balance:
@@ -599,6 +592,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 = staking_balance - existential_deposit
+        else:
+            staking_balance = staking_balance
+
     # Ask before moving on.
     if prompt:
         if not Confirm.ask(
@@ -615,7 +618,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 +1321,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 +1331,7 @@ async def delegate_stake(
         subtensor,
         wallet,
         delegate_ss58key,
-        Balance.from_tao(amount),
+        amount,
         wait_for_inclusion=True,
         prompt=prompt,
         delegate=True,
@@ -1338,7 +1341,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 +1351,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

@@ -26,6 +26,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
@@ -219,12 +222,12 @@ def test_root_commands(local_chain, wallet_setup):
     delegate_stake = Balance.from_tao(string_tao_to_float(alice_delegates_info[3]))
     assert delegate_stake == Balance.from_tao(delegate_amount)
 
-    # TOTAL STAKE(τ): This should be 10 as only Alice has delegated to Bob
+    # TOTAL STAKE(τ): This should be `delegate_amount` 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)
 
     # Total delegated Tao: This is listed at the bottom of the information
-    # Since Alice has only delegated to Bob, total should be 10 TAO
+    # Since Alice has only delegated to Bob, total should be `delegate_amount` TAO
     total_delegated_tao = Balance.from_tao(
         string_tao_to_float(alice_delegates.stdout.splitlines()[8].split()[3])
     )
@@ -257,6 +260,71 @@ def test_root_commands(local_chain, wallet_setup):
     )
     assert "✅ Finalized" in undelegate_alice.stdout
 
+    # 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)
+
+    # 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
+
+    # 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
+
+    # SS58: address of the Bob's hotkey (Alice has staked to Bob)
+    assert wallet_bob.hotkey.ss58_address == alice_delegates_info[2]
+
+    # Delegation: This should be 999999 as Alice delegated 999999 TAO to Bob
+    delegate_stake = Balance.from_tao(string_tao_to_float(alice_delegates_info[3]))
+    assert delegate_stake == Balance.from_tao(999999)
+
+    # 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="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",
+            "--all",
+            "--no-prompt",
+        ],
+    )
+    assert "✅ Finalized" in undelegate_alice.stdout
+
     print("✅ Passed Root commands")