Cleaned up logic, added tests

Author: thewhaleking

bittensor_cli/cli.py

@@ -388,18 +388,16 @@ def verbosity_console_handler(verbosity_level: int = 1) -> None:
         json_console.quiet = False
 
 
-def get_optional_netuid(
-    netuid: Optional[Union[int, list[int]]], all_netuids: bool
-) -> Optional[int]:
+def get_optional_netuid(netuid: Optional[int], all_netuids: bool) -> Optional[int]:
     """
     Parses options to determine if the user wants to use a specific netuid or all netuids (None)
 
     Returns:
         None if using all netuids, otherwise int for the netuid to use
     """
-    if (blank_netuid := (netuid is None or netuid == [])) and all_netuids is True:
+    if netuid is None and all_netuids is True:
         return None
-    elif blank_netuid and all_netuids is False:
+    elif netuid is None and all_netuids is False:
         answer = Prompt.ask(
             f"Enter the [{COLORS.G.SUBHEAD_MAIN}]netuid"
             f"[/{COLORS.G.SUBHEAD_MAIN}] to use. Leave blank for all netuids",
@@ -3276,29 +3274,42 @@ def stake_add(
 
         """
         netuids = netuids or []
-        if netuids:
-            netuids = parse_to_list(
-                netuids, int, "Netuids must be ints separated by commas", False
-            )
         self.verbosity_handler(quiet, verbose, json_output)
         safe_staking = self.ask_safe_staking(safe_staking)
         if safe_staking:
             rate_tolerance = self.ask_rate_tolerance(rate_tolerance)
             allow_partial_stake = self.ask_partial_stake(allow_partial_stake)
             console.print("\n")
-        if not netuids:
-            netuid_ = get_optional_netuid(netuids, all_netuids)
+
+        if netuids:
+            netuids = parse_to_list(
+                netuids, int, "Netuids must be ints separated by commas", False
+            )
+        else:
+            netuid_ = get_optional_netuid(None, all_netuids)
             netuids = [netuid_] if netuid_ else None
+        if netuids:
+            for netuid_ in netuids:
+                # ensure no negative netuids make it into our list
+                validate_netuid(netuid_)
 
         if stake_all and amount:
             print_error(
                 "Cannot specify an amount and 'stake-all'. Choose one or the other."
             )
-            raise typer.Exit()
+            return
 
         if stake_all and not amount:
             if not Confirm.ask("Stake all the available TAO tokens?", default=False):
-                raise typer.Exit()
+                return
+
+        if (
+            stake_all
+            and (isinstance(netuids, list) and len(netuids) > 1)
+            or (netuids is None)
+        ):
+            print_error("Cannot stake all to multiple subnets.")
+            return
 
         if all_hotkeys and include_hotkeys:
             print_error(

tests/e2e_tests/test_staking_sudo.py

@@ -36,6 +36,7 @@ def test_staking(local_chain, wallet_setup):
     """
     print("Testing staking and sudo commands🧪")
     netuid = 2
+    multiple_netuids = [2, 3]
     wallet_path_alice = "//Alice"
 
     # Create wallet for Alice
@@ -91,7 +92,42 @@ def test_staking(local_chain, wallet_setup):
     )
     result_output = json.loads(result.stdout)
     assert result_output["success"] is True
-    assert result_output["netuid"] == 2
+    assert result_output["netuid"] == netuid
+
+    # Register another subnet with sudo as Alice
+    result_for_second_repo = exec_command_alice(
+        command="subnets",
+        sub_command="create",
+        extra_args=[
+            "--wallet-path",
+            wallet_path_alice,
+            "--chain",
+            "ws://127.0.0.1:9945",
+            "--wallet-name",
+            wallet_alice.name,
+            "--wallet-hotkey",
+            wallet_alice.hotkey_str,
+            "--subnet-name",
+            "Test Subnet",
+            "--repo",
+            "https://github.com/username/repo",
+            "--contact",
+            "[email protected]",
+            "--url",
+            "https://testsubnet.com",
+            "--discord",
+            "alice#1234",
+            "--description",
+            "A test subnet for e2e testing",
+            "--additional-info",
+            "Created by Alice",
+            "--no-prompt",
+            "--json-output",
+        ],
+    )
+    result_output_second = json.loads(result_for_second_repo.stdout)
+    assert result_output_second["success"] is True
+    assert result_output_second["netuid"] == multiple_netuids[1]
 
     # Register Alice in netuid = 1 using her hotkey
     register_subnet = exec_command_alice(
@@ -192,7 +228,7 @@ def test_staking(local_chain, wallet_setup):
     assert get_identity_output["additional"] == sn_add_info
 
     # Add stake to Alice's hotkey
-    add_stake = exec_command_alice(
+    add_stake_single = exec_command_alice(
         command="stake",
         sub_command="add",
         extra_args=[
@@ -216,10 +252,10 @@ def test_staking(local_chain, wallet_setup):
             "144",
         ],
     )
-    assert "✅ Finalized" in add_stake.stdout, add_stake.stderr
+    assert "✅ Finalized" in add_stake_single.stdout, add_stake_single.stderr
 
     # Execute stake show for Alice's wallet
-    show_stake = exec_command_alice(
+    show_stake_adding_single = exec_command_alice(
         command="stake",
         sub_command="list",
         extra_args=[
@@ -235,7 +271,8 @@ def test_staking(local_chain, wallet_setup):
 
     # Assert correct stake is added
     cleaned_stake = [
-        re.sub(r"\s+", " ", line) for line in show_stake.stdout.splitlines()
+        re.sub(r"\s+", " ", line)
+        for line in show_stake_adding_single.stdout.splitlines()
     ]
     stake_added = cleaned_stake[8].split("│")[3].strip().split()[0]
     assert Balance.from_tao(float(stake_added)) >= Balance.from_tao(90)
@@ -284,6 +321,36 @@ def test_staking(local_chain, wallet_setup):
     )
     assert "✅ Finalized" in remove_stake.stdout
 
+    add_stake_multiple = exec_command_alice(
+        command="stake",
+        sub_command="add",
+        extra_args=[
+            "--netuids",
+            ",".join(str(x) for x in multiple_netuids),
+            "--wallet-path",
+            wallet_path_alice,
+            "--wallet-name",
+            wallet_alice.name,
+            "--hotkey",
+            wallet_alice.hotkey_str,
+            "--chain",
+            "ws://127.0.0.1:9945",
+            "--amount",
+            "100",
+            "--tolerance",
+            "0.1",
+            "--partial",
+            "--no-prompt",
+            "--era",
+            "144",
+        ],
+    )
+    assert "✅ Finalized" in add_stake_multiple.stdout, add_stake_multiple.stderr
+    for netuid_ in multiple_netuids:
+        assert f"Stake added to netuid: {netuid_}" in add_stake_multiple.stdout, (
+            add_stake_multiple.stderr
+        )
+
     # Fetch the hyperparameters of the subnet
     hyperparams = exec_command_alice(
         command="sudo",