From e819fb3da683c0906c9d9efaf2bdf72ae6db4d52 Mon Sep 17 00:00:00 2001 From: Mobile-Crest Date: Wed, 3 Dec 2025 09:47:47 +0100 Subject: [PATCH 1/6] feat: add cli for axon set and reset - Closes #225 --- bittensor_cli/cli.py | 168 +++++++ .../src/bittensor/extrinsics/serving.py | 214 ++++++++ tests/e2e_tests/test_axon.py | 372 ++++++++++++++ tests/unit_tests/test_axon_commands.py | 460 ++++++++++++++++++ 4 files changed, 1214 insertions(+) create mode 100644 bittensor_cli/src/bittensor/extrinsics/serving.py create mode 100644 tests/e2e_tests/test_axon.py create mode 100644 tests/unit_tests/test_axon_commands.py diff --git a/bittensor_cli/cli.py b/bittensor_cli/cli.py index 9ad6c89df..90ff8c86e 100755 --- a/bittensor_cli/cli.py +++ b/bittensor_cli/cli.py @@ -96,6 +96,7 @@ mechanisms as subnet_mechanisms, ) from bittensor_cli.src.commands.wallets import SortByBalance +from bittensor_cli.src.bittensor.extrinsics import serving from bittensor_cli.version import __version__, __version_as_int__ try: @@ -770,6 +771,7 @@ def __init__(self): self.liquidity_app = typer.Typer(epilog=_epilog) self.crowd_app = typer.Typer(epilog=_epilog) self.utils_app = typer.Typer(epilog=_epilog) + self.axon_app = typer.Typer(epilog=_epilog) # config alias self.app.add_typer( @@ -868,6 +870,14 @@ def __init__(self): no_args_is_help=True, ) + # axon app + self.app.add_typer( + self.axon_app, + name="axon", + short_help="Axon serving commands", + no_args_is_help=True, + ) + # config commands self.config_app.command("set")(self.set_config) self.config_app.command("get")(self.get_config) @@ -947,6 +957,10 @@ def __init__(self): "verify", rich_help_panel=HELP_PANELS["WALLET"]["OPERATIONS"] )(self.wallet_verify) + # axon commands + self.axon_app.command("reset")(self.axon_reset) + self.axon_app.command("set")(self.axon_set) + # stake commands self.stake_app.command( "add", rich_help_panel=HELP_PANELS["STAKE"]["STAKE_MGMT"] @@ -3717,6 +3731,160 @@ def wallet_swap_coldkey( ) ) + def axon_reset( + self, + netuid: int = Options.netuid, + wallet_name: Optional[str] = Options.wallet_name, + wallet_path: Optional[str] = Options.wallet_path, + wallet_hotkey: Optional[str] = Options.wallet_hotkey, + network: Optional[list[str]] = Options.network, + prompt: bool = Options.prompt, + wait_for_inclusion: bool = Options.wait_for_inclusion, + wait_for_finalization: bool = Options.wait_for_finalization, + quiet: bool = Options.quiet, + verbose: bool = Options.verbose, + ): + """ + Reset the axon information for a neuron on the network. + + This command removes the serving endpoint by setting the IP to 0.0.0.0 and port to 0, + indicating the neuron is no longer serving. + + USAGE + + The command requires you to specify the netuid where the neuron is registered. + It will reset the axon information for the hotkey associated with the wallet. + + EXAMPLE + + [green]$[/green] btcli axon reset --netuid 1 --wallet-name my_wallet --wallet-hotkey my_hotkey + + [bold]NOTE[/bold]: This command is used to stop serving on a specific subnet. The neuron will + remain registered but will not be reachable by other neurons until a new axon is set. + """ + self.verbosity_handler(quiet, verbose, False) + + wallet = self.wallet_ask( + wallet_name, + wallet_path, + wallet_hotkey, + ask_for=[WO.NAME, WO.HOTKEY], + validate=WV.WALLET_AND_HOTKEY, + ) + + subtensor = self.initialize_chain(network) + + logger.debug( + "args:\n" + f"netuid: {netuid}\n" + f"wallet: {wallet}\n" + f"prompt: {prompt}\n" + f"wait_for_inclusion: {wait_for_inclusion}\n" + f"wait_for_finalization: {wait_for_finalization}\n" + ) + + return self._run_command( + serving.reset_axon_extrinsic( + subtensor=subtensor, + wallet=wallet, + netuid=netuid, + prompt=prompt, + wait_for_inclusion=wait_for_inclusion, + wait_for_finalization=wait_for_finalization, + ) + ) + + def axon_set( + self, + netuid: int = Options.netuid, + ip: str = typer.Option( + ..., + "--ip", + help="IP address to set for the axon (e.g., '192.168.1.1')", + ), + port: int = typer.Option( + ..., + "--port", + help="Port number to set for the axon (0-65535)", + ), + ip_type: int = typer.Option( + 4, + "--ip-type", + help="IP type (4 for IPv4, 6 for IPv6)", + ), + protocol: int = typer.Option( + 4, + "--protocol", + help="Protocol version", + ), + wallet_name: Optional[str] = Options.wallet_name, + wallet_path: Optional[str] = Options.wallet_path, + wallet_hotkey: Optional[str] = Options.wallet_hotkey, + network: Optional[list[str]] = Options.network, + prompt: bool = Options.prompt, + wait_for_inclusion: bool = Options.wait_for_inclusion, + wait_for_finalization: bool = Options.wait_for_finalization, + quiet: bool = Options.quiet, + verbose: bool = Options.verbose, + ): + """ + Set the axon information for a neuron on the network. + + This command configures the serving endpoint for a neuron by specifying its IP address + and port, allowing other neurons to connect to it. + + USAGE + + The command requires you to specify the netuid, IP address, and port number. + It will set the axon information for the hotkey associated with the wallet. + + EXAMPLE + + [green]$[/green] btcli axon set --netuid 1 --ip 192.168.1.100 --port 8091 --wallet-name my_wallet --wallet-hotkey my_hotkey + + [bold]NOTE[/bold]: This command is used to advertise your serving endpoint on the network. + Make sure the IP and port are accessible from the internet if you want other neurons to connect. + """ + self.verbosity_handler(quiet, verbose, False) + + wallet = self.wallet_ask( + wallet_name, + wallet_path, + wallet_hotkey, + ask_for=[WO.NAME, WO.HOTKEY], + validate=WV.WALLET_AND_HOTKEY, + ) + + subtensor = self.initialize_chain(network) + + logger.debug( + "args:\n" + f"netuid: {netuid}\n" + f"ip: {ip}\n" + f"port: {port}\n" + f"ip_type: {ip_type}\n" + f"protocol: {protocol}\n" + f"wallet: {wallet}\n" + f"prompt: {prompt}\n" + f"wait_for_inclusion: {wait_for_inclusion}\n" + f"wait_for_finalization: {wait_for_finalization}\n" + ) + + return self._run_command( + serving.set_axon_extrinsic( + subtensor=subtensor, + wallet=wallet, + netuid=netuid, + ip=ip, + port=port, + ip_type=ip_type, + protocol=protocol, + prompt=prompt, + wait_for_inclusion=wait_for_inclusion, + wait_for_finalization=wait_for_finalization, + ) + ) + def get_auto_stake( self, network: Optional[list[str]] = Options.network, diff --git a/bittensor_cli/src/bittensor/extrinsics/serving.py b/bittensor_cli/src/bittensor/extrinsics/serving.py new file mode 100644 index 000000000..f94441ee1 --- /dev/null +++ b/bittensor_cli/src/bittensor/extrinsics/serving.py @@ -0,0 +1,214 @@ +""" +Extrinsics for serving operations (axon management). +""" +import typing +from typing import Optional + +from bittensor_wallet import Wallet +from rich.prompt import Confirm + +from bittensor_cli.src.bittensor.utils import ( + console, + err_console, + format_error_message, + unlock_key, + print_extrinsic_id, +) +from bittensor_cli.src.bittensor.networking import int_to_ip + +if typing.TYPE_CHECKING: + from bittensor_cli.src.bittensor.subtensor_interface import SubtensorInterface + + +def ip_to_int(ip_str: str) -> int: + """ + Converts an IP address string to its integer representation. + + Args: + ip_str: IP address string (e.g., "192.168.1.1") + + Returns: + Integer representation of the IP address + """ + import netaddr + return int(netaddr.IPAddress(ip_str)) + + +async def reset_axon_extrinsic( + subtensor: "SubtensorInterface", + wallet: Wallet, + netuid: int, + prompt: bool = False, + wait_for_inclusion: bool = True, + wait_for_finalization: bool = False, +) -> tuple[bool, str]: + """ + Resets the axon information for a neuron on the network. + + This effectively removes the serving endpoint by setting the IP to 0.0.0.0 + and port to 0, indicating the neuron is no longer serving. + + Args: + subtensor: The subtensor interface to use for the extrinsic + wallet: The wallet containing the hotkey to reset the axon for + netuid: The network UID where the neuron is registered + prompt: Whether to prompt for confirmation before submitting + wait_for_inclusion: Whether to wait for the extrinsic to be included in a block + wait_for_finalization: Whether to wait for the extrinsic to be finalized + + Returns: + Tuple of (success: bool, message: str) + """ + # Unlock the hotkey + if not unlock_key(wallet, hotkey=True).success: + return False, "Failed to unlock hotkey" + + # Prompt for confirmation if requested + if prompt: + if not Confirm.ask( + f"Do you want to reset the axon for hotkey [bold]{wallet.hotkey.ss58_address}[/bold] " + f"on netuid [bold]{netuid}[/bold]?" + ): + return False, "User cancelled the operation" + + with console.status( + f":satellite: Resetting axon on [white]netuid {netuid}[/white]..." + ): + try: + # Compose the serve_axon call with reset values (IP: 0.0.0.0, port: 0) + call = await subtensor.substrate.compose_call( + call_module="SubtensorModule", + call_function="serve_axon", + call_params={ + "netuid": netuid, + "version": 0, + "ip": ip_to_int("0.0.0.0"), + "port": 0, + "ip_type": 4, # IPv4 + "protocol": 4, + "placeholder1": 0, + "placeholder2": 0, + }, + ) + + # Sign and submit the extrinsic + success, error_msg, response = await subtensor.sign_and_send_extrinsic( + call=call, + wallet=wallet, + wait_for_inclusion=wait_for_inclusion, + wait_for_finalization=wait_for_finalization, + ) + + if success: + if wait_for_inclusion or wait_for_finalization: + await print_extrinsic_id(response) + console.print(":white_heavy_check_mark: [green]Axon reset successfully[/green]") + return True, "Axon reset successfully" + else: + err_console.print(f":cross_mark: [red]Failed to reset axon: {error_msg}[/red]") + return False, error_msg + + except Exception as e: + error_message = format_error_message(e) + err_console.print(f":cross_mark: [red]Failed to reset axon: {error_message}[/red]") + return False, error_message + + +async def set_axon_extrinsic( + subtensor: "SubtensorInterface", + wallet: Wallet, + netuid: int, + ip: str, + port: int, + ip_type: int = 4, + protocol: int = 4, + prompt: bool = False, + wait_for_inclusion: bool = True, + wait_for_finalization: bool = False, +) -> tuple[bool, str]: + """ + Sets the axon information for a neuron on the network. + + This configures the serving endpoint for a neuron by specifying its IP address + and port, allowing other neurons to connect to it. + + Args: + subtensor: The subtensor interface to use for the extrinsic + wallet: The wallet containing the hotkey to set the axon for + netuid: The network UID where the neuron is registered + ip: The IP address to set (e.g., "192.168.1.1") + port: The port number to set + ip_type: IP type (4 for IPv4, 6 for IPv6) + protocol: Protocol version (default: 4) + prompt: Whether to prompt for confirmation before submitting + wait_for_inclusion: Whether to wait for the extrinsic to be included in a block + wait_for_finalization: Whether to wait for the extrinsic to be finalized + + Returns: + Tuple of (success: bool, message: str) + """ + # Validate port + if not (0 <= port <= 65535): + return False, f"Invalid port number: {port}. Must be between 0 and 65535." + + # Validate IP address + try: + ip_int = ip_to_int(ip) + except Exception as e: + return False, f"Invalid IP address: {ip}. Error: {str(e)}" + + # Unlock the hotkey + if not unlock_key(wallet, hotkey=True).success: + return False, "Failed to unlock hotkey" + + # Prompt for confirmation if requested + if prompt: + if not Confirm.ask( + f"Do you want to set the axon for hotkey [bold]{wallet.hotkey.ss58_address}[/bold] " + f"on netuid [bold]{netuid}[/bold] to [bold]{ip}:{port}[/bold]?" + ): + return False, "User cancelled the operation" + + with console.status( + f":satellite: Setting axon on [white]netuid {netuid}[/white] to [white]{ip}:{port}[/white]..." + ): + try: + # Compose the serve_axon call + call = await subtensor.substrate.compose_call( + call_module="SubtensorModule", + call_function="serve_axon", + call_params={ + "netuid": netuid, + "version": 0, + "ip": ip_int, + "port": port, + "ip_type": ip_type, + "protocol": protocol, + "placeholder1": 0, + "placeholder2": 0, + }, + ) + + # Sign and submit the extrinsic + success, error_msg, response = await subtensor.sign_and_send_extrinsic( + call=call, + wallet=wallet, + wait_for_inclusion=wait_for_inclusion, + wait_for_finalization=wait_for_finalization, + ) + + if success: + if wait_for_inclusion or wait_for_finalization: + await print_extrinsic_id(response) + console.print( + f":white_heavy_check_mark: [green]Axon set successfully to {ip}:{port}[/green]" + ) + return True, f"Axon set successfully to {ip}:{port}" + else: + err_console.print(f":cross_mark: [red]Failed to set axon: {error_msg}[/red]") + return False, error_msg + + except Exception as e: + error_message = format_error_message(e) + err_console.print(f":cross_mark: [red]Failed to set axon: {error_message}[/red]") + return False, error_message diff --git a/tests/e2e_tests/test_axon.py b/tests/e2e_tests/test_axon.py new file mode 100644 index 000000000..6e79a68df --- /dev/null +++ b/tests/e2e_tests/test_axon.py @@ -0,0 +1,372 @@ +""" +End-to-end tests for axon commands. + +Verify commands: +* btcli axon reset +* btcli axon set +""" +import pytest +import re + + +@pytest.mark.parametrize("local_chain", [None], indirect=True) +def test_axon_reset_and_set(local_chain, wallet_setup): + """ + Test axon reset and set commands end-to-end. + + This test: + 1. Creates a subnet + 2. Registers a neuron + 3. Sets the axon information + 4. Verifies the axon is set correctly + 5. Resets the axon + 6. Verifies the axon is reset (0.0.0.0:0) + """ + wallet_path_alice = "//Alice" + netuid = 1 + + # Create wallet for Alice + keypair_alice, wallet_alice, wallet_path_alice, exec_command_alice = wallet_setup( + wallet_path_alice + ) + + # Register a subnet with sudo as Alice + result = 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, + "--no-prompt", + ], + ) + assert result.exit_code == 0, f"Subnet creation failed: {result.stdout}" + + # Register neuron on the subnet + result = exec_command_alice( + command="subnets", + sub_command="register", + 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, + "--netuid", + str(netuid), + "--no-prompt", + ], + ) + assert result.exit_code == 0, f"Neuron registration failed: {result.stdout}" + + # Set axon information + test_ip = "192.168.1.100" + test_port = 8091 + + result = exec_command_alice( + command="axon", + sub_command="set", + 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, + "--netuid", + str(netuid), + "--ip", + test_ip, + "--port", + str(test_port), + "--no-prompt", + ], + ) + + assert result.exit_code == 0, f"Axon set failed: {result.stdout}" + assert "successfully" in result.stdout.lower() or "success" in result.stdout.lower(), \ + f"Success message not found in output: {result.stdout}" + + # Verify axon is set by checking wallet overview + result = exec_command_alice( + command="wallet", + sub_command="overview", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_alice.name, + "--netuid", + str(netuid), + ], + ) + + assert result.exit_code == 0, f"Wallet overview failed: {result.stdout}" + + # Check that axon column shows an IP (not "none") + # The overview should show the axon info in the AXON column + lines = result.stdout.split("\n") + axon_found = False + for line in lines: + # Look for a line with the neuron info that has an IP address in the AXON column + if wallet_alice.hotkey_str[:8] in line and "none" not in line.lower(): + # Check if there's an IP-like pattern in the line + if re.search(r"\d+\.\d+\.\d+\.\d+:\d+", line): + axon_found = True + break + + assert axon_found, f"Axon not set correctly in overview: {result.stdout}" + + # Reset axon + result = exec_command_alice( + command="axon", + sub_command="reset", + 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, + "--netuid", + str(netuid), + "--no-prompt", + ], + ) + + assert result.exit_code == 0, f"Axon reset failed: {result.stdout}" + assert "successfully" in result.stdout.lower() or "success" in result.stdout.lower(), \ + f"Success message not found in output: {result.stdout}" + + # Verify axon is reset by checking wallet overview + result = exec_command_alice( + command="wallet", + sub_command="overview", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_alice.name, + "--netuid", + str(netuid), + ], + ) + + assert result.exit_code == 0, f"Wallet overview failed: {result.stdout}" + + # Check that axon column shows "none" after reset + lines = result.stdout.split("\n") + axon_reset = False + for line in lines: + if wallet_alice.hotkey_str[:8] in line and "none" in line.lower(): + axon_reset = True + break + + assert axon_reset, f"Axon not reset correctly in overview: {result.stdout}" + + +@pytest.mark.parametrize("local_chain", [None], indirect=True) +def test_axon_set_with_ipv6(local_chain, wallet_setup): + """ + Test setting axon with IPv6 address. + """ + wallet_path_bob = "//Bob" + netuid = 1 + + # Create wallet for Bob + keypair_bob, wallet_bob, wallet_path_bob, exec_command_bob = wallet_setup( + wallet_path_bob + ) + + # Register a subnet with sudo as Bob + result = exec_command_bob( + command="subnets", + sub_command="create", + extra_args=[ + "--wallet-path", + wallet_path_bob, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_bob.name, + "--wallet-hotkey", + wallet_bob.hotkey_str, + "--no-prompt", + ], + ) + assert result.exit_code == 0, f"Subnet creation failed: {result.stdout}" + + # Register neuron on the subnet + result = exec_command_bob( + command="subnets", + sub_command="register", + extra_args=[ + "--wallet-path", + wallet_path_bob, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_bob.name, + "--wallet-hotkey", + wallet_bob.hotkey_str, + "--netuid", + str(netuid), + "--no-prompt", + ], + ) + assert result.exit_code == 0, f"Neuron registration failed: {result.stdout}" + + # Set axon with IPv6 address + test_ipv6 = "2001:db8::1" + test_port = 8092 + + result = exec_command_bob( + command="axon", + sub_command="set", + extra_args=[ + "--wallet-path", + wallet_path_bob, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_bob.name, + "--wallet-hotkey", + wallet_bob.hotkey_str, + "--netuid", + str(netuid), + "--ip", + test_ipv6, + "--port", + str(test_port), + "--ip-type", + "6", # IPv6 + "--no-prompt", + ], + ) + + assert result.exit_code == 0, f"Axon set with IPv6 failed: {result.stdout}" + assert "successfully" in result.stdout.lower() or "success" in result.stdout.lower(), \ + f"Success message not found in output: {result.stdout}" + + +@pytest.mark.parametrize("local_chain", [None], indirect=True) +def test_axon_set_invalid_inputs(local_chain, wallet_setup): + """ + Test axon set with invalid inputs to ensure proper error handling. + """ + wallet_path_charlie = "//Charlie" + netuid = 1 + + # Create wallet for Charlie + keypair_charlie, wallet_charlie, wallet_path_charlie, exec_command_charlie = wallet_setup( + wallet_path_charlie + ) + + # Register a subnet + result = exec_command_charlie( + command="subnets", + sub_command="create", + extra_args=[ + "--wallet-path", + wallet_path_charlie, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_charlie.name, + "--wallet-hotkey", + wallet_charlie.hotkey_str, + "--no-prompt", + ], + ) + assert result.exit_code == 0 + + # Register neuron + result = exec_command_charlie( + command="subnets", + sub_command="register", + extra_args=[ + "--wallet-path", + wallet_path_charlie, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_charlie.name, + "--wallet-hotkey", + wallet_charlie.hotkey_str, + "--netuid", + str(netuid), + "--no-prompt", + ], + ) + assert result.exit_code == 0 + + # Test with invalid port (too high) + result = exec_command_charlie( + command="axon", + sub_command="set", + extra_args=[ + "--wallet-path", + wallet_path_charlie, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_charlie.name, + "--wallet-hotkey", + wallet_charlie.hotkey_str, + "--netuid", + str(netuid), + "--ip", + "192.168.1.1", + "--port", + "70000", # Invalid port + "--no-prompt", + ], + ) + + # Should fail with invalid port + assert result.exit_code != 0 or "invalid port" in result.stdout.lower() or "failed" in result.stdout.lower(), \ + f"Expected error for invalid port, got: {result.stdout}" + + # Test with invalid IP + result = exec_command_charlie( + command="axon", + sub_command="set", + extra_args=[ + "--wallet-path", + wallet_path_charlie, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_charlie.name, + "--wallet-hotkey", + wallet_charlie.hotkey_str, + "--netuid", + str(netuid), + "--ip", + "invalid.ip.address", # Invalid IP + "--port", + "8091", + "--no-prompt", + ], + ) + + # Should fail with invalid IP + assert result.exit_code != 0 or "invalid ip" in result.stdout.lower() or "failed" in result.stdout.lower(), \ + f"Expected error for invalid IP, got: {result.stdout}" diff --git a/tests/unit_tests/test_axon_commands.py b/tests/unit_tests/test_axon_commands.py new file mode 100644 index 000000000..49ef24134 --- /dev/null +++ b/tests/unit_tests/test_axon_commands.py @@ -0,0 +1,460 @@ +""" +Unit tests for axon commands (reset and set). +""" +import pytest +from unittest.mock import AsyncMock, MagicMock, Mock, patch +from bittensor_wallet import Wallet + +from bittensor_cli.src.bittensor.extrinsics.serving import ( + reset_axon_extrinsic, + set_axon_extrinsic, + ip_to_int, +) + + +class TestIpToInt: + """Tests for IP address to integer conversion.""" + + def test_ipv4_conversion(self): + """Test IPv4 address conversion.""" + assert ip_to_int("0.0.0.0") == 0 + assert ip_to_int("127.0.0.1") == 2130706433 + assert ip_to_int("192.168.1.1") == 3232235777 + assert ip_to_int("255.255.255.255") == 4294967295 + + def test_ipv6_conversion(self): + """Test IPv6 address conversion.""" + # IPv6 loopback + result = ip_to_int("::1") + assert result == 1 + + # IPv6 address + result = ip_to_int("2001:db8::1") + assert result > 0 + + def test_invalid_ip_raises_error(self): + """Test that invalid IP addresses raise errors.""" + with pytest.raises(Exception): + ip_to_int("invalid.ip.address") + + with pytest.raises(Exception): + ip_to_int("256.256.256.256") + + +class TestResetAxonExtrinsic: + """Tests for reset_axon_extrinsic function.""" + + @pytest.mark.asyncio + async def test_reset_axon_success(self): + """Test successful axon reset.""" + # Setup mocks + mock_subtensor = MagicMock() + mock_subtensor.substrate.compose_call = AsyncMock(return_value="mock_call") + mock_response = MagicMock() + mock_response.is_success = AsyncMock(return_value=True) + mock_response.get_extrinsic_identifier = AsyncMock(return_value="0x123") + mock_subtensor.sign_and_send_extrinsic = AsyncMock( + return_value=(True, "", mock_response) + ) + + mock_wallet = MagicMock(spec=Wallet) + mock_wallet.hotkey.ss58_address = "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + + with ( + patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock, + patch("bittensor_cli.src.bittensor.extrinsics.serving.print_extrinsic_id", new_callable=AsyncMock), + ): + mock_unlock.return_value = MagicMock(success=True) + + # Execute + success, message = await reset_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + prompt=False, + wait_for_inclusion=True, + wait_for_finalization=False, + ) + + # Verify + assert success is True + assert "successfully" in message.lower() + + # Verify compose_call was called with correct parameters + mock_subtensor.substrate.compose_call.assert_called_once() + call_args = mock_subtensor.substrate.compose_call.call_args + assert call_args[1]["call_module"] == "SubtensorModule" + assert call_args[1]["call_function"] == "serve_axon" + assert call_args[1]["call_params"]["netuid"] == 1 + assert call_args[1]["call_params"]["ip"] == 0 # 0.0.0.0 as int + assert call_args[1]["call_params"]["port"] == 0 + assert call_args[1]["call_params"]["ip_type"] == 4 + + @pytest.mark.asyncio + async def test_reset_axon_unlock_failure(self): + """Test axon reset when hotkey unlock fails.""" + mock_subtensor = MagicMock() + mock_wallet = MagicMock(spec=Wallet) + + with patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock: + mock_unlock.return_value = MagicMock(success=False) + + success, message = await reset_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + prompt=False, + ) + + assert success is False + assert "unlock" in message.lower() + + @pytest.mark.asyncio + async def test_reset_axon_user_cancellation(self): + """Test axon reset when user cancels prompt.""" + mock_subtensor = MagicMock() + mock_wallet = MagicMock(spec=Wallet) + mock_wallet.hotkey.ss58_address = "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + + with ( + patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock, + patch("bittensor_cli.src.bittensor.extrinsics.serving.Confirm") as mock_confirm, + ): + mock_unlock.return_value = MagicMock(success=True) + mock_confirm.ask.return_value = False + + success, message = await reset_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + prompt=True, + ) + + assert success is False + assert "cancelled" in message.lower() + + @pytest.mark.asyncio + async def test_reset_axon_extrinsic_failure(self): + """Test axon reset when extrinsic submission fails.""" + mock_subtensor = MagicMock() + mock_subtensor.substrate.compose_call = AsyncMock(return_value="mock_call") + mock_subtensor.sign_and_send_extrinsic = AsyncMock( + return_value=(False, "Network error", None) + ) + + mock_wallet = MagicMock(spec=Wallet) + mock_wallet.hotkey.ss58_address = "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + + with patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock: + mock_unlock.return_value = MagicMock(success=True) + + success, message = await reset_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + prompt=False, + ) + + assert success is False + assert "Network error" in message + + +class TestSetAxonExtrinsic: + """Tests for set_axon_extrinsic function.""" + + @pytest.mark.asyncio + async def test_set_axon_success(self): + """Test successful axon set.""" + mock_subtensor = MagicMock() + mock_subtensor.substrate.compose_call = AsyncMock(return_value="mock_call") + mock_response = MagicMock() + mock_response.is_success = AsyncMock(return_value=True) + mock_response.get_extrinsic_identifier = AsyncMock(return_value="0x123") + mock_subtensor.sign_and_send_extrinsic = AsyncMock( + return_value=(True, "", mock_response) + ) + + mock_wallet = MagicMock(spec=Wallet) + mock_wallet.hotkey.ss58_address = "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + + with ( + patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock, + patch("bittensor_cli.src.bittensor.extrinsics.serving.print_extrinsic_id", new_callable=AsyncMock), + ): + mock_unlock.return_value = MagicMock(success=True) + + success, message = await set_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + ip="192.168.1.100", + port=8091, + ip_type=4, + protocol=4, + prompt=False, + wait_for_inclusion=True, + wait_for_finalization=False, + ) + + assert success is True + assert "successfully" in message.lower() + assert "192.168.1.100:8091" in message + + # Verify compose_call was called with correct parameters + mock_subtensor.substrate.compose_call.assert_called_once() + call_args = mock_subtensor.substrate.compose_call.call_args + assert call_args[1]["call_module"] == "SubtensorModule" + assert call_args[1]["call_function"] == "serve_axon" + assert call_args[1]["call_params"]["netuid"] == 1 + assert call_args[1]["call_params"]["port"] == 8091 + assert call_args[1]["call_params"]["ip_type"] == 4 + assert call_args[1]["call_params"]["protocol"] == 4 + + @pytest.mark.asyncio + async def test_set_axon_invalid_port(self): + """Test axon set with invalid port number.""" + mock_subtensor = MagicMock() + mock_wallet = MagicMock(spec=Wallet) + + # Test port too high + success, message = await set_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + ip="192.168.1.100", + port=70000, + prompt=False, + ) + + assert success is False + assert "Invalid port" in message + + # Test negative port + success, message = await set_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + ip="192.168.1.100", + port=-1, + prompt=False, + ) + + assert success is False + assert "Invalid port" in message + + @pytest.mark.asyncio + async def test_set_axon_invalid_ip(self): + """Test axon set with invalid IP address.""" + mock_subtensor = MagicMock() + mock_wallet = MagicMock(spec=Wallet) + + success, message = await set_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + ip="invalid.ip.address", + port=8091, + prompt=False, + ) + + assert success is False + assert "Invalid IP" in message + + @pytest.mark.asyncio + async def test_set_axon_unlock_failure(self): + """Test axon set when hotkey unlock fails.""" + mock_subtensor = MagicMock() + mock_wallet = MagicMock(spec=Wallet) + + with patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock: + mock_unlock.return_value = MagicMock(success=False) + + success, message = await set_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + ip="192.168.1.100", + port=8091, + prompt=False, + ) + + assert success is False + assert "unlock" in message.lower() + + @pytest.mark.asyncio + async def test_set_axon_user_cancellation(self): + """Test axon set when user cancels prompt.""" + mock_subtensor = MagicMock() + mock_wallet = MagicMock(spec=Wallet) + mock_wallet.hotkey.ss58_address = "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + + with ( + patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock, + patch("bittensor_cli.src.bittensor.extrinsics.serving.Confirm") as mock_confirm, + ): + mock_unlock.return_value = MagicMock(success=True) + mock_confirm.ask.return_value = False + + success, message = await set_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + ip="192.168.1.100", + port=8091, + prompt=True, + ) + + assert success is False + assert "cancelled" in message.lower() + + @pytest.mark.asyncio + async def test_set_axon_with_ipv6(self): + """Test axon set with IPv6 address.""" + mock_subtensor = MagicMock() + mock_subtensor.substrate.compose_call = AsyncMock(return_value="mock_call") + mock_response = MagicMock() + mock_response.is_success = AsyncMock(return_value=True) + mock_response.get_extrinsic_identifier = AsyncMock(return_value="0x123") + mock_subtensor.sign_and_send_extrinsic = AsyncMock( + return_value=(True, "", mock_response) + ) + + mock_wallet = MagicMock(spec=Wallet) + mock_wallet.hotkey.ss58_address = "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + + with ( + patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock, + patch("bittensor_cli.src.bittensor.extrinsics.serving.print_extrinsic_id", new_callable=AsyncMock), + ): + mock_unlock.return_value = MagicMock(success=True) + + success, message = await set_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + ip="2001:db8::1", + port=8091, + ip_type=6, + protocol=4, + prompt=False, + ) + + assert success is True + + # Verify ip_type was set to 6 + call_args = mock_subtensor.substrate.compose_call.call_args + assert call_args[1]["call_params"]["ip_type"] == 6 + + @pytest.mark.asyncio + async def test_set_axon_exception_handling(self): + """Test axon set handles exceptions gracefully.""" + mock_subtensor = MagicMock() + mock_subtensor.substrate.compose_call = AsyncMock( + side_effect=Exception("Unexpected error") + ) + + mock_wallet = MagicMock(spec=Wallet) + mock_wallet.hotkey.ss58_address = "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + + with patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock: + mock_unlock.return_value = MagicMock(success=True) + + success, message = await set_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + ip="192.168.1.100", + port=8091, + prompt=False, + ) + + assert success is False + assert len(message) > 0 + + +class TestAxonCLICommands: + """Tests for CLI command handlers.""" + + @patch("bittensor_cli.cli.serving") + def test_axon_reset_command_handler(self, mock_serving): + """Test axon reset CLI command handler.""" + from bittensor_cli.cli import CLIManager + + cli_manager = CLIManager() + mock_serving.reset_axon_extrinsic = AsyncMock( + return_value=(True, "Success") + ) + + with ( + patch.object(cli_manager, "verbosity_handler"), + patch.object(cli_manager, "wallet_ask") as mock_wallet_ask, + patch.object(cli_manager, "initialize_chain") as mock_init_chain, + patch.object(cli_manager, "_run_command") as mock_run_command, + ): + mock_wallet = Mock() + mock_wallet_ask.return_value = mock_wallet + mock_subtensor = Mock() + mock_init_chain.return_value = mock_subtensor + + cli_manager.axon_reset( + netuid=1, + wallet_name="test_wallet", + wallet_path="/tmp/test", + wallet_hotkey="test_hotkey", + network=None, + prompt=False, + wait_for_inclusion=True, + wait_for_finalization=False, + quiet=False, + verbose=False, + ) + + # Verify wallet_ask was called correctly + mock_wallet_ask.assert_called_once() + + # Verify _run_command was called + mock_run_command.assert_called_once() + + @patch("bittensor_cli.cli.serving") + def test_axon_set_command_handler(self, mock_serving): + """Test axon set CLI command handler.""" + from bittensor_cli.cli import CLIManager + + cli_manager = CLIManager() + mock_serving.set_axon_extrinsic = AsyncMock( + return_value=(True, "Success") + ) + + with ( + patch.object(cli_manager, "verbosity_handler"), + patch.object(cli_manager, "wallet_ask") as mock_wallet_ask, + patch.object(cli_manager, "initialize_chain") as mock_init_chain, + patch.object(cli_manager, "_run_command") as mock_run_command, + ): + mock_wallet = Mock() + mock_wallet_ask.return_value = mock_wallet + mock_subtensor = Mock() + mock_init_chain.return_value = mock_subtensor + + cli_manager.axon_set( + netuid=1, + ip="192.168.1.100", + port=8091, + ip_type=4, + protocol=4, + wallet_name="test_wallet", + wallet_path="/tmp/test", + wallet_hotkey="test_hotkey", + network=None, + prompt=False, + wait_for_inclusion=True, + wait_for_finalization=False, + quiet=False, + verbose=False, + ) + + # Verify wallet_ask was called correctly + mock_wallet_ask.assert_called_once() + + # Verify _run_command was called + mock_run_command.assert_called_once() From 763ee28c4f23f8d790683e2cfbdb392107da6464 Mon Sep 17 00:00:00 2001 From: Mobile-Crest Date: Wed, 3 Dec 2025 17:57:13 +0100 Subject: [PATCH 2/6] refactoring and bug fix --- bittensor_cli/cli.py | 22 +-- .../src/bittensor/extrinsics/serving.py | 93 ++++++++----- bittensor_cli/src/commands/axon/__init__.py | 0 bittensor_cli/src/commands/axon/axon.py | 128 ++++++++++++++++++ 4 files changed, 202 insertions(+), 41 deletions(-) create mode 100644 bittensor_cli/src/commands/axon/__init__.py create mode 100644 bittensor_cli/src/commands/axon/axon.py diff --git a/bittensor_cli/cli.py b/bittensor_cli/cli.py index 90ff8c86e..dbcd6d2ef 100755 --- a/bittensor_cli/cli.py +++ b/bittensor_cli/cli.py @@ -95,8 +95,8 @@ subnets, mechanisms as subnet_mechanisms, ) +from bittensor_cli.src.commands.axon import axon from bittensor_cli.src.commands.wallets import SortByBalance -from bittensor_cli.src.bittensor.extrinsics import serving from bittensor_cli.version import __version__, __version_as_int__ try: @@ -3743,11 +3743,12 @@ def axon_reset( wait_for_finalization: bool = Options.wait_for_finalization, quiet: bool = Options.quiet, verbose: bool = Options.verbose, + json_output: bool = Options.json_output, ): """ Reset the axon information for a neuron on the network. - This command removes the serving endpoint by setting the IP to 0.0.0.0 and port to 0, + This command removes the serving endpoint by setting the IP to 0.0.0.0 and port to 1, indicating the neuron is no longer serving. USAGE @@ -3762,7 +3763,7 @@ def axon_reset( [bold]NOTE[/bold]: This command is used to stop serving on a specific subnet. The neuron will remain registered but will not be reachable by other neurons until a new axon is set. """ - self.verbosity_handler(quiet, verbose, False) + self.verbosity_handler(quiet, verbose, json_output) wallet = self.wallet_ask( wallet_name, @@ -3784,13 +3785,14 @@ def axon_reset( ) return self._run_command( - serving.reset_axon_extrinsic( - subtensor=subtensor, + axon.reset( wallet=wallet, + subtensor=subtensor, netuid=netuid, prompt=prompt, wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, + json_output=json_output, ) ) @@ -3801,11 +3803,13 @@ def axon_set( ..., "--ip", help="IP address to set for the axon (e.g., '192.168.1.1')", + prompt="Enter the IP address for the axon (e.g., '192.168.1.1' or '2001:db8::1')", ), port: int = typer.Option( ..., "--port", help="Port number to set for the axon (0-65535)", + prompt="Enter the port number for the axon (0-65535)", ), ip_type: int = typer.Option( 4, @@ -3826,6 +3830,7 @@ def axon_set( wait_for_finalization: bool = Options.wait_for_finalization, quiet: bool = Options.quiet, verbose: bool = Options.verbose, + json_output: bool = Options.json_output, ): """ Set the axon information for a neuron on the network. @@ -3845,7 +3850,7 @@ def axon_set( [bold]NOTE[/bold]: This command is used to advertise your serving endpoint on the network. Make sure the IP and port are accessible from the internet if you want other neurons to connect. """ - self.verbosity_handler(quiet, verbose, False) + self.verbosity_handler(quiet, verbose, json_output) wallet = self.wallet_ask( wallet_name, @@ -3871,9 +3876,9 @@ def axon_set( ) return self._run_command( - serving.set_axon_extrinsic( - subtensor=subtensor, + axon.set_axon( wallet=wallet, + subtensor=subtensor, netuid=netuid, ip=ip, port=port, @@ -3882,6 +3887,7 @@ def axon_set( prompt=prompt, wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, + json_output=json_output, ) ) diff --git a/bittensor_cli/src/bittensor/extrinsics/serving.py b/bittensor_cli/src/bittensor/extrinsics/serving.py index f94441ee1..f67e9ad87 100644 --- a/bittensor_cli/src/bittensor/extrinsics/serving.py +++ b/bittensor_cli/src/bittensor/extrinsics/serving.py @@ -41,7 +41,7 @@ async def reset_axon_extrinsic( prompt: bool = False, wait_for_inclusion: bool = True, wait_for_finalization: bool = False, -) -> tuple[bool, str]: +) -> tuple[bool, str, Optional[str]]: """ Resets the axon information for a neuron on the network. @@ -57,11 +57,11 @@ async def reset_axon_extrinsic( wait_for_finalization: Whether to wait for the extrinsic to be finalized Returns: - Tuple of (success: bool, message: str) + Tuple of (success: bool, message: str, extrinsic_id: Optional[str]) """ # Unlock the hotkey - if not unlock_key(wallet, hotkey=True).success: - return False, "Failed to unlock hotkey" + if not (unlock_status := unlock_key(wallet, unlock_type="hot", print_out=False)).success: + return False, unlock_status.message, None # Prompt for confirmation if requested if prompt: @@ -69,13 +69,14 @@ async def reset_axon_extrinsic( f"Do you want to reset the axon for hotkey [bold]{wallet.hotkey.ss58_address}[/bold] " f"on netuid [bold]{netuid}[/bold]?" ): - return False, "User cancelled the operation" + return False, "User cancelled the operation", None with console.status( f":satellite: Resetting axon on [white]netuid {netuid}[/white]..." ): try: - # Compose the serve_axon call with reset values (IP: 0.0.0.0, port: 0) + # Compose the serve_axon call with reset values (IP: 0.0.0.0, port: 1) + # Note: Port must be >= 1 as chain rejects port 0 as invalid call = await subtensor.substrate.compose_call( call_module="SubtensorModule", call_function="serve_axon", @@ -83,7 +84,7 @@ async def reset_axon_extrinsic( "netuid": netuid, "version": 0, "ip": ip_to_int("0.0.0.0"), - "port": 0, + "port": 1, "ip_type": 4, # IPv4 "protocol": 4, "placeholder1": 0, @@ -91,27 +92,41 @@ async def reset_axon_extrinsic( }, ) - # Sign and submit the extrinsic - success, error_msg, response = await subtensor.sign_and_send_extrinsic( + # Sign with hotkey and submit the extrinsic + extrinsic = await subtensor.substrate.create_signed_extrinsic( call=call, - wallet=wallet, + keypair=wallet.hotkey, + ) + response = await subtensor.substrate.submit_extrinsic( + extrinsic, wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, ) - if success: - if wait_for_inclusion or wait_for_finalization: - await print_extrinsic_id(response) - console.print(":white_heavy_check_mark: [green]Axon reset successfully[/green]") - return True, "Axon reset successfully" + # We only wait here if we expect finalization. + if not wait_for_finalization and not wait_for_inclusion: + console.print( + ":white_heavy_check_mark: [dark_sea_green3]Axon reset successfully[/dark_sea_green3]" + ) + return True, "Not waiting for finalization or inclusion.", None + + success = await response.is_success + if not success: + error_msg = format_error_message(await response.error_message) + err_console.print(f":cross_mark: [red]Failed[/red]: {error_msg}") + return False, error_msg, None else: - err_console.print(f":cross_mark: [red]Failed to reset axon: {error_msg}[/red]") - return False, error_msg + ext_id = await response.get_extrinsic_identifier() + await print_extrinsic_id(response) + console.print( + ":white_heavy_check_mark: [dark_sea_green3]Axon reset successfully[/dark_sea_green3]" + ) + return True, "Axon reset successfully", ext_id except Exception as e: error_message = format_error_message(e) err_console.print(f":cross_mark: [red]Failed to reset axon: {error_message}[/red]") - return False, error_message + return False, error_message, None async def set_axon_extrinsic( @@ -125,7 +140,7 @@ async def set_axon_extrinsic( prompt: bool = False, wait_for_inclusion: bool = True, wait_for_finalization: bool = False, -) -> tuple[bool, str]: +) -> tuple[bool, str, Optional[str]]: """ Sets the axon information for a neuron on the network. @@ -158,8 +173,8 @@ async def set_axon_extrinsic( return False, f"Invalid IP address: {ip}. Error: {str(e)}" # Unlock the hotkey - if not unlock_key(wallet, hotkey=True).success: - return False, "Failed to unlock hotkey" + if not (unlock_status := unlock_key(wallet, unlock_type="hot", print_out=False)).success: + return False, unlock_status.message # Prompt for confirmation if requested if prompt: @@ -167,7 +182,7 @@ async def set_axon_extrinsic( f"Do you want to set the axon for hotkey [bold]{wallet.hotkey.ss58_address}[/bold] " f"on netuid [bold]{netuid}[/bold] to [bold]{ip}:{port}[/bold]?" ): - return False, "User cancelled the operation" + return False, "User cancelled the operation", None with console.status( f":satellite: Setting axon on [white]netuid {netuid}[/white] to [white]{ip}:{port}[/white]..." @@ -189,26 +204,38 @@ async def set_axon_extrinsic( }, ) - # Sign and submit the extrinsic - success, error_msg, response = await subtensor.sign_and_send_extrinsic( + # Sign with hotkey and submit the extrinsic + extrinsic = await subtensor.substrate.create_signed_extrinsic( call=call, - wallet=wallet, + keypair=wallet.hotkey, + ) + response = await subtensor.substrate.submit_extrinsic( + extrinsic, wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, ) - if success: - if wait_for_inclusion or wait_for_finalization: - await print_extrinsic_id(response) + # We only wait here if we expect finalization. + if not wait_for_finalization and not wait_for_inclusion: console.print( - f":white_heavy_check_mark: [green]Axon set successfully to {ip}:{port}[/green]" + f":white_heavy_check_mark: [dark_sea_green3]Axon set successfully to {ip}:{port}[/dark_sea_green3]" ) - return True, f"Axon set successfully to {ip}:{port}" + return True, "Not waiting for finalization or inclusion.", None + + success = await response.is_success + if not success: + error_msg = format_error_message(await response.error_message) + err_console.print(f":cross_mark: [red]Failed[/red]: {error_msg}") + return False, error_msg, None else: - err_console.print(f":cross_mark: [red]Failed to set axon: {error_msg}[/red]") - return False, error_msg + ext_id = await response.get_extrinsic_identifier() + await print_extrinsic_id(response) + console.print( + f":white_heavy_check_mark: [dark_sea_green3]Axon set successfully to {ip}:{port}[/dark_sea_green3]" + ) + return True, f"Axon set successfully to {ip}:{port}", ext_id except Exception as e: error_message = format_error_message(e) err_console.print(f":cross_mark: [red]Failed to set axon: {error_message}[/red]") - return False, error_message + return False, error_message, None diff --git a/bittensor_cli/src/commands/axon/__init__.py b/bittensor_cli/src/commands/axon/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/bittensor_cli/src/commands/axon/axon.py b/bittensor_cli/src/commands/axon/axon.py new file mode 100644 index 000000000..ef600f3ed --- /dev/null +++ b/bittensor_cli/src/commands/axon/axon.py @@ -0,0 +1,128 @@ +""" +Axon commands for managing neuron serving endpoints. +""" +import json +from typing import TYPE_CHECKING + +from bittensor_wallet import Wallet + +from bittensor_cli.src.bittensor.utils import ( + console, + err_console, + json_console, +) +from bittensor_cli.src.bittensor.extrinsics.serving import ( + reset_axon_extrinsic, + set_axon_extrinsic, +) + +if TYPE_CHECKING: + from bittensor_cli.src.bittensor.subtensor_interface import SubtensorInterface + + +async def reset( + wallet: Wallet, + subtensor: "SubtensorInterface", + netuid: int, + json_output: bool, + prompt: bool, + wait_for_inclusion: bool, + wait_for_finalization: bool, +): + """ + Reset the axon information for a neuron on the network. + + This command removes the serving endpoint by setting the IP to 0.0.0.0 and port to 1, + indicating the neuron is no longer serving. + + Args: + wallet: The wallet containing the hotkey to reset the axon for + subtensor: The subtensor interface to use for the extrinsic + netuid: The network UID where the neuron is registered + json_output: Whether to output results in JSON format + prompt: Whether to prompt for confirmation before submitting + wait_for_inclusion: Whether to wait for the extrinsic to be included in a block + wait_for_finalization: Whether to wait for the extrinsic to be finalized + """ + success, message, ext_id = await reset_axon_extrinsic( + subtensor=subtensor, + wallet=wallet, + netuid=netuid, + prompt=prompt, + wait_for_inclusion=wait_for_inclusion, + wait_for_finalization=wait_for_finalization, + ) + + if json_output: + json_console.print( + json.dumps({ + "success": success, + "message": message, + "extrinsic_identifier": ext_id, + "netuid": netuid, + "hotkey": wallet.hotkey.ss58_address, + }) + ) + elif not success: + err_console.print(f"[red]Failed to reset axon: {message}[/red]") + + +async def set_axon( + wallet: Wallet, + subtensor: "SubtensorInterface", + netuid: int, + ip: str, + port: int, + ip_type: int, + protocol: int, + json_output: bool, + prompt: bool, + wait_for_inclusion: bool, + wait_for_finalization: bool, +): + """ + Set the axon information for a neuron on the network. + + This command configures the serving endpoint for a neuron by specifying its IP address + and port, allowing other neurons to connect to it. + + Args: + wallet: The wallet containing the hotkey to set the axon for + subtensor: The subtensor interface to use for the extrinsic + netuid: The network UID where the neuron is registered + ip: IP address to set for the axon + port: Port number to set for the axon + ip_type: IP type (4 for IPv4, 6 for IPv6) + protocol: Protocol version + json_output: Whether to output results in JSON format + prompt: Whether to prompt for confirmation before submitting + wait_for_inclusion: Whether to wait for the extrinsic to be included in a block + wait_for_finalization: Whether to wait for the extrinsic to be finalized + """ + success, message, ext_id = await set_axon_extrinsic( + subtensor=subtensor, + wallet=wallet, + netuid=netuid, + ip=ip, + port=port, + ip_type=ip_type, + protocol=protocol, + prompt=prompt, + wait_for_inclusion=wait_for_inclusion, + wait_for_finalization=wait_for_finalization, + ) + + if json_output: + json_console.print( + json.dumps({ + "success": success, + "message": message, + "extrinsic_identifier": ext_id, + "netuid": netuid, + "hotkey": wallet.hotkey.ss58_address, + "ip": ip, + "port": port, + }) + ) + elif not success: + err_console.print(f"[red]Failed to set axon: {message}[/red]") From 8e50f2b22ef2f06be85fb8f3322dc8225d376c38 Mon Sep 17 00:00:00 2001 From: Mobile-Crest Date: Wed, 3 Dec 2025 20:41:57 +0100 Subject: [PATCH 3/6] fix unit and e2e tests --- .../src/bittensor/extrinsics/serving.py | 8 +- tests/e2e_tests/test_axon.py | 2 +- tests/unit_tests/test_axon_commands.py | 101 +++++++++++------- 3 files changed, 66 insertions(+), 45 deletions(-) diff --git a/bittensor_cli/src/bittensor/extrinsics/serving.py b/bittensor_cli/src/bittensor/extrinsics/serving.py index f67e9ad87..e1762ede4 100644 --- a/bittensor_cli/src/bittensor/extrinsics/serving.py +++ b/bittensor_cli/src/bittensor/extrinsics/serving.py @@ -160,21 +160,21 @@ async def set_axon_extrinsic( wait_for_finalization: Whether to wait for the extrinsic to be finalized Returns: - Tuple of (success: bool, message: str) + Tuple of (success: bool, message: str, extrinsic_id: Optional[str]) """ # Validate port if not (0 <= port <= 65535): - return False, f"Invalid port number: {port}. Must be between 0 and 65535." + return False, f"Invalid port number: {port}. Must be between 0 and 65535.", None # Validate IP address try: ip_int = ip_to_int(ip) except Exception as e: - return False, f"Invalid IP address: {ip}. Error: {str(e)}" + return False, f"Invalid IP address: {ip}. Error: {str(e)}", None # Unlock the hotkey if not (unlock_status := unlock_key(wallet, unlock_type="hot", print_out=False)).success: - return False, unlock_status.message + return False, unlock_status.message, None # Prompt for confirmation if requested if prompt: diff --git a/tests/e2e_tests/test_axon.py b/tests/e2e_tests/test_axon.py index 6e79a68df..ba50812c7 100644 --- a/tests/e2e_tests/test_axon.py +++ b/tests/e2e_tests/test_axon.py @@ -20,7 +20,7 @@ def test_axon_reset_and_set(local_chain, wallet_setup): 3. Sets the axon information 4. Verifies the axon is set correctly 5. Resets the axon - 6. Verifies the axon is reset (0.0.0.0:0) + 6. Verifies the axon is reset (0.0.0.0:1 - not serving) """ wallet_path_alice = "//Alice" netuid = 1 diff --git a/tests/unit_tests/test_axon_commands.py b/tests/unit_tests/test_axon_commands.py index 49ef24134..4ac4e8ec8 100644 --- a/tests/unit_tests/test_axon_commands.py +++ b/tests/unit_tests/test_axon_commands.py @@ -50,12 +50,14 @@ async def test_reset_axon_success(self): # Setup mocks mock_subtensor = MagicMock() mock_subtensor.substrate.compose_call = AsyncMock(return_value="mock_call") + mock_subtensor.substrate.create_signed_extrinsic = AsyncMock(return_value="mock_extrinsic") mock_response = MagicMock() - mock_response.is_success = AsyncMock(return_value=True) + # is_success is a property that returns a coroutine + async def mock_is_success(): + return True + mock_response.is_success = mock_is_success() mock_response.get_extrinsic_identifier = AsyncMock(return_value="0x123") - mock_subtensor.sign_and_send_extrinsic = AsyncMock( - return_value=(True, "", mock_response) - ) + mock_subtensor.substrate.submit_extrinsic = AsyncMock(return_value=mock_response) mock_wallet = MagicMock(spec=Wallet) mock_wallet.hotkey.ss58_address = "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" @@ -67,7 +69,7 @@ async def test_reset_axon_success(self): mock_unlock.return_value = MagicMock(success=True) # Execute - success, message = await reset_axon_extrinsic( + success, message, ext_id = await reset_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, netuid=1, @@ -79,6 +81,7 @@ async def test_reset_axon_success(self): # Verify assert success is True assert "successfully" in message.lower() + assert ext_id == "0x123" # Verify compose_call was called with correct parameters mock_subtensor.substrate.compose_call.assert_called_once() @@ -87,7 +90,7 @@ async def test_reset_axon_success(self): assert call_args[1]["call_function"] == "serve_axon" assert call_args[1]["call_params"]["netuid"] == 1 assert call_args[1]["call_params"]["ip"] == 0 # 0.0.0.0 as int - assert call_args[1]["call_params"]["port"] == 0 + assert call_args[1]["call_params"]["port"] == 1 # Port 1, not 0 assert call_args[1]["call_params"]["ip_type"] == 4 @pytest.mark.asyncio @@ -97,9 +100,9 @@ async def test_reset_axon_unlock_failure(self): mock_wallet = MagicMock(spec=Wallet) with patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock: - mock_unlock.return_value = MagicMock(success=False) + mock_unlock.return_value = MagicMock(success=False, message="Failed to unlock hotkey") - success, message = await reset_axon_extrinsic( + success, message, ext_id = await reset_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, netuid=1, @@ -108,6 +111,7 @@ async def test_reset_axon_unlock_failure(self): assert success is False assert "unlock" in message.lower() + assert ext_id is None @pytest.mark.asyncio async def test_reset_axon_user_cancellation(self): @@ -123,7 +127,7 @@ async def test_reset_axon_user_cancellation(self): mock_unlock.return_value = MagicMock(success=True) mock_confirm.ask.return_value = False - success, message = await reset_axon_extrinsic( + success, message, ext_id = await reset_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, netuid=1, @@ -132,15 +136,20 @@ async def test_reset_axon_user_cancellation(self): assert success is False assert "cancelled" in message.lower() + assert ext_id is None @pytest.mark.asyncio async def test_reset_axon_extrinsic_failure(self): """Test axon reset when extrinsic submission fails.""" mock_subtensor = MagicMock() mock_subtensor.substrate.compose_call = AsyncMock(return_value="mock_call") - mock_subtensor.sign_and_send_extrinsic = AsyncMock( - return_value=(False, "Network error", None) - ) + mock_subtensor.substrate.create_signed_extrinsic = AsyncMock(return_value="mock_extrinsic") + mock_response = MagicMock() + async def mock_is_success(): + return False + mock_response.is_success = mock_is_success() + mock_response.error_message = AsyncMock(return_value="Network error") + mock_subtensor.substrate.submit_extrinsic = AsyncMock(return_value=mock_response) mock_wallet = MagicMock(spec=Wallet) mock_wallet.hotkey.ss58_address = "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" @@ -148,7 +157,7 @@ async def test_reset_axon_extrinsic_failure(self): with patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock: mock_unlock.return_value = MagicMock(success=True) - success, message = await reset_axon_extrinsic( + success, message, ext_id = await reset_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, netuid=1, @@ -156,7 +165,8 @@ async def test_reset_axon_extrinsic_failure(self): ) assert success is False - assert "Network error" in message + assert len(message) > 0 + assert ext_id is None class TestSetAxonExtrinsic: @@ -167,12 +177,13 @@ async def test_set_axon_success(self): """Test successful axon set.""" mock_subtensor = MagicMock() mock_subtensor.substrate.compose_call = AsyncMock(return_value="mock_call") + mock_subtensor.substrate.create_signed_extrinsic = AsyncMock(return_value="mock_extrinsic") mock_response = MagicMock() - mock_response.is_success = AsyncMock(return_value=True) + async def mock_is_success(): + return True + mock_response.is_success = mock_is_success() mock_response.get_extrinsic_identifier = AsyncMock(return_value="0x123") - mock_subtensor.sign_and_send_extrinsic = AsyncMock( - return_value=(True, "", mock_response) - ) + mock_subtensor.substrate.submit_extrinsic = AsyncMock(return_value=mock_response) mock_wallet = MagicMock(spec=Wallet) mock_wallet.hotkey.ss58_address = "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" @@ -183,7 +194,7 @@ async def test_set_axon_success(self): ): mock_unlock.return_value = MagicMock(success=True) - success, message = await set_axon_extrinsic( + success, message, ext_id = await set_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, netuid=1, @@ -198,6 +209,7 @@ async def test_set_axon_success(self): assert success is True assert "successfully" in message.lower() + assert ext_id == "0x123" assert "192.168.1.100:8091" in message # Verify compose_call was called with correct parameters @@ -217,7 +229,7 @@ async def test_set_axon_invalid_port(self): mock_wallet = MagicMock(spec=Wallet) # Test port too high - success, message = await set_axon_extrinsic( + success, message, ext_id = await set_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, netuid=1, @@ -228,9 +240,10 @@ async def test_set_axon_invalid_port(self): assert success is False assert "Invalid port" in message + assert ext_id is None # Test negative port - success, message = await set_axon_extrinsic( + success, message, ext_id = await set_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, netuid=1, @@ -241,6 +254,7 @@ async def test_set_axon_invalid_port(self): assert success is False assert "Invalid port" in message + assert ext_id is None @pytest.mark.asyncio async def test_set_axon_invalid_ip(self): @@ -248,7 +262,7 @@ async def test_set_axon_invalid_ip(self): mock_subtensor = MagicMock() mock_wallet = MagicMock(spec=Wallet) - success, message = await set_axon_extrinsic( + success, message, ext_id = await set_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, netuid=1, @@ -259,6 +273,7 @@ async def test_set_axon_invalid_ip(self): assert success is False assert "Invalid IP" in message + assert ext_id is None @pytest.mark.asyncio async def test_set_axon_unlock_failure(self): @@ -267,9 +282,9 @@ async def test_set_axon_unlock_failure(self): mock_wallet = MagicMock(spec=Wallet) with patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock: - mock_unlock.return_value = MagicMock(success=False) + mock_unlock.return_value = MagicMock(success=False, message="Failed to unlock hotkey") - success, message = await set_axon_extrinsic( + success, message, ext_id = await set_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, netuid=1, @@ -280,6 +295,7 @@ async def test_set_axon_unlock_failure(self): assert success is False assert "unlock" in message.lower() + assert "Failed to unlock hotkey" in message @pytest.mark.asyncio async def test_set_axon_user_cancellation(self): @@ -295,7 +311,7 @@ async def test_set_axon_user_cancellation(self): mock_unlock.return_value = MagicMock(success=True) mock_confirm.ask.return_value = False - success, message = await set_axon_extrinsic( + success, message, ext_id = await set_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, netuid=1, @@ -312,12 +328,13 @@ async def test_set_axon_with_ipv6(self): """Test axon set with IPv6 address.""" mock_subtensor = MagicMock() mock_subtensor.substrate.compose_call = AsyncMock(return_value="mock_call") + mock_subtensor.substrate.create_signed_extrinsic = AsyncMock(return_value="mock_extrinsic") mock_response = MagicMock() - mock_response.is_success = AsyncMock(return_value=True) + async def mock_is_success(): + return True + mock_response.is_success = mock_is_success() mock_response.get_extrinsic_identifier = AsyncMock(return_value="0x123") - mock_subtensor.sign_and_send_extrinsic = AsyncMock( - return_value=(True, "", mock_response) - ) + mock_subtensor.substrate.submit_extrinsic = AsyncMock(return_value=mock_response) mock_wallet = MagicMock(spec=Wallet) mock_wallet.hotkey.ss58_address = "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" @@ -328,7 +345,7 @@ async def test_set_axon_with_ipv6(self): ): mock_unlock.return_value = MagicMock(success=True) - success, message = await set_axon_extrinsic( + success, message, ext_id = await set_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, netuid=1, @@ -340,7 +357,8 @@ async def test_set_axon_with_ipv6(self): ) assert success is True - + assert "successfully" in message.lower() + assert ext_id == "0x123" # Verify ip_type was set to 6 call_args = mock_subtensor.substrate.compose_call.call_args assert call_args[1]["call_params"]["ip_type"] == 6 @@ -359,7 +377,7 @@ async def test_set_axon_exception_handling(self): with patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock: mock_unlock.return_value = MagicMock(success=True) - success, message = await set_axon_extrinsic( + success, message, ext_id = await set_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, netuid=1, @@ -370,19 +388,20 @@ async def test_set_axon_exception_handling(self): assert success is False assert len(message) > 0 + assert ext_id is None class TestAxonCLICommands: """Tests for CLI command handlers.""" - @patch("bittensor_cli.cli.serving") - def test_axon_reset_command_handler(self, mock_serving): + @patch("bittensor_cli.cli.axon") + def test_axon_reset_command_handler(self, mock_axon): """Test axon reset CLI command handler.""" from bittensor_cli.cli import CLIManager cli_manager = CLIManager() - mock_serving.reset_axon_extrinsic = AsyncMock( - return_value=(True, "Success") + mock_axon.reset = AsyncMock( + return_value=None ) with ( @@ -407,6 +426,7 @@ def test_axon_reset_command_handler(self, mock_serving): wait_for_finalization=False, quiet=False, verbose=False, + json_output=False, ) # Verify wallet_ask was called correctly @@ -415,14 +435,14 @@ def test_axon_reset_command_handler(self, mock_serving): # Verify _run_command was called mock_run_command.assert_called_once() - @patch("bittensor_cli.cli.serving") - def test_axon_set_command_handler(self, mock_serving): + @patch("bittensor_cli.cli.axon") + def test_axon_set_command_handler(self, mock_axon): """Test axon set CLI command handler.""" from bittensor_cli.cli import CLIManager cli_manager = CLIManager() - mock_serving.set_axon_extrinsic = AsyncMock( - return_value=(True, "Success") + mock_axon.set_axon = AsyncMock( + return_value=None ) with ( @@ -451,6 +471,7 @@ def test_axon_set_command_handler(self, mock_serving): wait_for_finalization=False, quiet=False, verbose=False, + json_output=False, ) # Verify wallet_ask was called correctly From b34eb123bc0525229c26964117e4cbca40933f26 Mon Sep 17 00:00:00 2001 From: Mobile-Crest Date: Fri, 5 Dec 2025 03:49:33 +0100 Subject: [PATCH 4/6] run ruff format --- bittensor_cli/cli.py | 40 +-- .../src/bittensor/extrinsics/serving.py | 62 +++-- bittensor_cli/src/commands/axon/axon.py | 49 ++-- tests/e2e_tests/test_axon.py | 94 ++++--- tests/unit_tests/test_axon_commands.py | 248 +++++++++++------- 5 files changed, 294 insertions(+), 199 deletions(-) diff --git a/bittensor_cli/cli.py b/bittensor_cli/cli.py index dbcd6d2ef..b6d91f8bb 100755 --- a/bittensor_cli/cli.py +++ b/bittensor_cli/cli.py @@ -3747,24 +3747,24 @@ def axon_reset( ): """ Reset the axon information for a neuron on the network. - + This command removes the serving endpoint by setting the IP to 0.0.0.0 and port to 1, indicating the neuron is no longer serving. - + USAGE - + The command requires you to specify the netuid where the neuron is registered. It will reset the axon information for the hotkey associated with the wallet. - + EXAMPLE - + [green]$[/green] btcli axon reset --netuid 1 --wallet-name my_wallet --wallet-hotkey my_hotkey - + [bold]NOTE[/bold]: This command is used to stop serving on a specific subnet. The neuron will remain registered but will not be reachable by other neurons until a new axon is set. """ self.verbosity_handler(quiet, verbose, json_output) - + wallet = self.wallet_ask( wallet_name, wallet_path, @@ -3772,9 +3772,9 @@ def axon_reset( ask_for=[WO.NAME, WO.HOTKEY], validate=WV.WALLET_AND_HOTKEY, ) - + subtensor = self.initialize_chain(network) - + logger.debug( "args:\n" f"netuid: {netuid}\n" @@ -3783,7 +3783,7 @@ def axon_reset( f"wait_for_inclusion: {wait_for_inclusion}\n" f"wait_for_finalization: {wait_for_finalization}\n" ) - + return self._run_command( axon.reset( wallet=wallet, @@ -3834,24 +3834,24 @@ def axon_set( ): """ Set the axon information for a neuron on the network. - + This command configures the serving endpoint for a neuron by specifying its IP address and port, allowing other neurons to connect to it. - + USAGE - + The command requires you to specify the netuid, IP address, and port number. It will set the axon information for the hotkey associated with the wallet. - + EXAMPLE - + [green]$[/green] btcli axon set --netuid 1 --ip 192.168.1.100 --port 8091 --wallet-name my_wallet --wallet-hotkey my_hotkey - + [bold]NOTE[/bold]: This command is used to advertise your serving endpoint on the network. Make sure the IP and port are accessible from the internet if you want other neurons to connect. """ self.verbosity_handler(quiet, verbose, json_output) - + wallet = self.wallet_ask( wallet_name, wallet_path, @@ -3859,9 +3859,9 @@ def axon_set( ask_for=[WO.NAME, WO.HOTKEY], validate=WV.WALLET_AND_HOTKEY, ) - + subtensor = self.initialize_chain(network) - + logger.debug( "args:\n" f"netuid: {netuid}\n" @@ -3874,7 +3874,7 @@ def axon_set( f"wait_for_inclusion: {wait_for_inclusion}\n" f"wait_for_finalization: {wait_for_finalization}\n" ) - + return self._run_command( axon.set_axon( wallet=wallet, diff --git a/bittensor_cli/src/bittensor/extrinsics/serving.py b/bittensor_cli/src/bittensor/extrinsics/serving.py index e1762ede4..4e88b6944 100644 --- a/bittensor_cli/src/bittensor/extrinsics/serving.py +++ b/bittensor_cli/src/bittensor/extrinsics/serving.py @@ -1,6 +1,7 @@ """ Extrinsics for serving operations (axon management). """ + import typing from typing import Optional @@ -23,14 +24,15 @@ def ip_to_int(ip_str: str) -> int: """ Converts an IP address string to its integer representation. - + Args: ip_str: IP address string (e.g., "192.168.1.1") - + Returns: Integer representation of the IP address """ import netaddr + return int(netaddr.IPAddress(ip_str)) @@ -44,10 +46,10 @@ async def reset_axon_extrinsic( ) -> tuple[bool, str, Optional[str]]: """ Resets the axon information for a neuron on the network. - + This effectively removes the serving endpoint by setting the IP to 0.0.0.0 and port to 0, indicating the neuron is no longer serving. - + Args: subtensor: The subtensor interface to use for the extrinsic wallet: The wallet containing the hotkey to reset the axon for @@ -55,14 +57,16 @@ async def reset_axon_extrinsic( prompt: Whether to prompt for confirmation before submitting wait_for_inclusion: Whether to wait for the extrinsic to be included in a block wait_for_finalization: Whether to wait for the extrinsic to be finalized - + Returns: Tuple of (success: bool, message: str, extrinsic_id: Optional[str]) """ # Unlock the hotkey - if not (unlock_status := unlock_key(wallet, unlock_type="hot", print_out=False)).success: + if not ( + unlock_status := unlock_key(wallet, unlock_type="hot", print_out=False) + ).success: return False, unlock_status.message, None - + # Prompt for confirmation if requested if prompt: if not Confirm.ask( @@ -70,7 +74,7 @@ async def reset_axon_extrinsic( f"on netuid [bold]{netuid}[/bold]?" ): return False, "User cancelled the operation", None - + with console.status( f":satellite: Resetting axon on [white]netuid {netuid}[/white]..." ): @@ -91,7 +95,7 @@ async def reset_axon_extrinsic( "placeholder2": 0, }, ) - + # Sign with hotkey and submit the extrinsic extrinsic = await subtensor.substrate.create_signed_extrinsic( call=call, @@ -102,14 +106,14 @@ async def reset_axon_extrinsic( wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, ) - + # We only wait here if we expect finalization. if not wait_for_finalization and not wait_for_inclusion: console.print( ":white_heavy_check_mark: [dark_sea_green3]Axon reset successfully[/dark_sea_green3]" ) return True, "Not waiting for finalization or inclusion.", None - + success = await response.is_success if not success: error_msg = format_error_message(await response.error_message) @@ -122,10 +126,12 @@ async def reset_axon_extrinsic( ":white_heavy_check_mark: [dark_sea_green3]Axon reset successfully[/dark_sea_green3]" ) return True, "Axon reset successfully", ext_id - + except Exception as e: error_message = format_error_message(e) - err_console.print(f":cross_mark: [red]Failed to reset axon: {error_message}[/red]") + err_console.print( + f":cross_mark: [red]Failed to reset axon: {error_message}[/red]" + ) return False, error_message, None @@ -143,10 +149,10 @@ async def set_axon_extrinsic( ) -> tuple[bool, str, Optional[str]]: """ Sets the axon information for a neuron on the network. - + This configures the serving endpoint for a neuron by specifying its IP address and port, allowing other neurons to connect to it. - + Args: subtensor: The subtensor interface to use for the extrinsic wallet: The wallet containing the hotkey to set the axon for @@ -158,24 +164,26 @@ async def set_axon_extrinsic( prompt: Whether to prompt for confirmation before submitting wait_for_inclusion: Whether to wait for the extrinsic to be included in a block wait_for_finalization: Whether to wait for the extrinsic to be finalized - + Returns: Tuple of (success: bool, message: str, extrinsic_id: Optional[str]) """ # Validate port if not (0 <= port <= 65535): return False, f"Invalid port number: {port}. Must be between 0 and 65535.", None - + # Validate IP address try: ip_int = ip_to_int(ip) except Exception as e: return False, f"Invalid IP address: {ip}. Error: {str(e)}", None - + # Unlock the hotkey - if not (unlock_status := unlock_key(wallet, unlock_type="hot", print_out=False)).success: + if not ( + unlock_status := unlock_key(wallet, unlock_type="hot", print_out=False) + ).success: return False, unlock_status.message, None - + # Prompt for confirmation if requested if prompt: if not Confirm.ask( @@ -183,7 +191,7 @@ async def set_axon_extrinsic( f"on netuid [bold]{netuid}[/bold] to [bold]{ip}:{port}[/bold]?" ): return False, "User cancelled the operation", None - + with console.status( f":satellite: Setting axon on [white]netuid {netuid}[/white] to [white]{ip}:{port}[/white]..." ): @@ -203,7 +211,7 @@ async def set_axon_extrinsic( "placeholder2": 0, }, ) - + # Sign with hotkey and submit the extrinsic extrinsic = await subtensor.substrate.create_signed_extrinsic( call=call, @@ -214,14 +222,14 @@ async def set_axon_extrinsic( wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, ) - + # We only wait here if we expect finalization. if not wait_for_finalization and not wait_for_inclusion: console.print( f":white_heavy_check_mark: [dark_sea_green3]Axon set successfully to {ip}:{port}[/dark_sea_green3]" ) return True, "Not waiting for finalization or inclusion.", None - + success = await response.is_success if not success: error_msg = format_error_message(await response.error_message) @@ -234,8 +242,10 @@ async def set_axon_extrinsic( f":white_heavy_check_mark: [dark_sea_green3]Axon set successfully to {ip}:{port}[/dark_sea_green3]" ) return True, f"Axon set successfully to {ip}:{port}", ext_id - + except Exception as e: error_message = format_error_message(e) - err_console.print(f":cross_mark: [red]Failed to set axon: {error_message}[/red]") + err_console.print( + f":cross_mark: [red]Failed to set axon: {error_message}[/red]" + ) return False, error_message, None diff --git a/bittensor_cli/src/commands/axon/axon.py b/bittensor_cli/src/commands/axon/axon.py index ef600f3ed..95db9134f 100644 --- a/bittensor_cli/src/commands/axon/axon.py +++ b/bittensor_cli/src/commands/axon/axon.py @@ -1,6 +1,7 @@ """ Axon commands for managing neuron serving endpoints. """ + import json from typing import TYPE_CHECKING @@ -31,10 +32,10 @@ async def reset( ): """ Reset the axon information for a neuron on the network. - + This command removes the serving endpoint by setting the IP to 0.0.0.0 and port to 1, indicating the neuron is no longer serving. - + Args: wallet: The wallet containing the hotkey to reset the axon for subtensor: The subtensor interface to use for the extrinsic @@ -52,16 +53,18 @@ async def reset( wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, ) - + if json_output: json_console.print( - json.dumps({ - "success": success, - "message": message, - "extrinsic_identifier": ext_id, - "netuid": netuid, - "hotkey": wallet.hotkey.ss58_address, - }) + json.dumps( + { + "success": success, + "message": message, + "extrinsic_identifier": ext_id, + "netuid": netuid, + "hotkey": wallet.hotkey.ss58_address, + } + ) ) elif not success: err_console.print(f"[red]Failed to reset axon: {message}[/red]") @@ -82,10 +85,10 @@ async def set_axon( ): """ Set the axon information for a neuron on the network. - + This command configures the serving endpoint for a neuron by specifying its IP address and port, allowing other neurons to connect to it. - + Args: wallet: The wallet containing the hotkey to set the axon for subtensor: The subtensor interface to use for the extrinsic @@ -111,18 +114,20 @@ async def set_axon( wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, ) - + if json_output: json_console.print( - json.dumps({ - "success": success, - "message": message, - "extrinsic_identifier": ext_id, - "netuid": netuid, - "hotkey": wallet.hotkey.ss58_address, - "ip": ip, - "port": port, - }) + json.dumps( + { + "success": success, + "message": message, + "extrinsic_identifier": ext_id, + "netuid": netuid, + "hotkey": wallet.hotkey.ss58_address, + "ip": ip, + "port": port, + } + ) ) elif not success: err_console.print(f"[red]Failed to set axon: {message}[/red]") diff --git a/tests/e2e_tests/test_axon.py b/tests/e2e_tests/test_axon.py index ba50812c7..5b3bf539f 100644 --- a/tests/e2e_tests/test_axon.py +++ b/tests/e2e_tests/test_axon.py @@ -5,6 +5,7 @@ * btcli axon reset * btcli axon set """ + import pytest import re @@ -13,7 +14,7 @@ def test_axon_reset_and_set(local_chain, wallet_setup): """ Test axon reset and set commands end-to-end. - + This test: 1. Creates a subnet 2. Registers a neuron @@ -24,12 +25,12 @@ def test_axon_reset_and_set(local_chain, wallet_setup): """ wallet_path_alice = "//Alice" netuid = 1 - + # Create wallet for Alice keypair_alice, wallet_alice, wallet_path_alice, exec_command_alice = wallet_setup( wallet_path_alice ) - + # Register a subnet with sudo as Alice result = exec_command_alice( command="subnets", @@ -47,7 +48,7 @@ def test_axon_reset_and_set(local_chain, wallet_setup): ], ) assert result.exit_code == 0, f"Subnet creation failed: {result.stdout}" - + # Register neuron on the subnet result = exec_command_alice( command="subnets", @@ -67,11 +68,11 @@ def test_axon_reset_and_set(local_chain, wallet_setup): ], ) assert result.exit_code == 0, f"Neuron registration failed: {result.stdout}" - + # Set axon information test_ip = "192.168.1.100" test_port = 8091 - + result = exec_command_alice( command="axon", sub_command="set", @@ -93,11 +94,12 @@ def test_axon_reset_and_set(local_chain, wallet_setup): "--no-prompt", ], ) - + assert result.exit_code == 0, f"Axon set failed: {result.stdout}" - assert "successfully" in result.stdout.lower() or "success" in result.stdout.lower(), \ - f"Success message not found in output: {result.stdout}" - + assert ( + "successfully" in result.stdout.lower() or "success" in result.stdout.lower() + ), f"Success message not found in output: {result.stdout}" + # Verify axon is set by checking wallet overview result = exec_command_alice( command="wallet", @@ -113,9 +115,9 @@ def test_axon_reset_and_set(local_chain, wallet_setup): str(netuid), ], ) - + assert result.exit_code == 0, f"Wallet overview failed: {result.stdout}" - + # Check that axon column shows an IP (not "none") # The overview should show the axon info in the AXON column lines = result.stdout.split("\n") @@ -127,9 +129,9 @@ def test_axon_reset_and_set(local_chain, wallet_setup): if re.search(r"\d+\.\d+\.\d+\.\d+:\d+", line): axon_found = True break - + assert axon_found, f"Axon not set correctly in overview: {result.stdout}" - + # Reset axon result = exec_command_alice( command="axon", @@ -148,11 +150,12 @@ def test_axon_reset_and_set(local_chain, wallet_setup): "--no-prompt", ], ) - + assert result.exit_code == 0, f"Axon reset failed: {result.stdout}" - assert "successfully" in result.stdout.lower() or "success" in result.stdout.lower(), \ - f"Success message not found in output: {result.stdout}" - + assert ( + "successfully" in result.stdout.lower() or "success" in result.stdout.lower() + ), f"Success message not found in output: {result.stdout}" + # Verify axon is reset by checking wallet overview result = exec_command_alice( command="wallet", @@ -168,9 +171,9 @@ def test_axon_reset_and_set(local_chain, wallet_setup): str(netuid), ], ) - + assert result.exit_code == 0, f"Wallet overview failed: {result.stdout}" - + # Check that axon column shows "none" after reset lines = result.stdout.split("\n") axon_reset = False @@ -178,7 +181,7 @@ def test_axon_reset_and_set(local_chain, wallet_setup): if wallet_alice.hotkey_str[:8] in line and "none" in line.lower(): axon_reset = True break - + assert axon_reset, f"Axon not reset correctly in overview: {result.stdout}" @@ -189,12 +192,12 @@ def test_axon_set_with_ipv6(local_chain, wallet_setup): """ wallet_path_bob = "//Bob" netuid = 1 - + # Create wallet for Bob keypair_bob, wallet_bob, wallet_path_bob, exec_command_bob = wallet_setup( wallet_path_bob ) - + # Register a subnet with sudo as Bob result = exec_command_bob( command="subnets", @@ -212,7 +215,7 @@ def test_axon_set_with_ipv6(local_chain, wallet_setup): ], ) assert result.exit_code == 0, f"Subnet creation failed: {result.stdout}" - + # Register neuron on the subnet result = exec_command_bob( command="subnets", @@ -232,11 +235,11 @@ def test_axon_set_with_ipv6(local_chain, wallet_setup): ], ) assert result.exit_code == 0, f"Neuron registration failed: {result.stdout}" - + # Set axon with IPv6 address test_ipv6 = "2001:db8::1" test_port = 8092 - + result = exec_command_bob( command="axon", sub_command="set", @@ -260,10 +263,11 @@ def test_axon_set_with_ipv6(local_chain, wallet_setup): "--no-prompt", ], ) - + assert result.exit_code == 0, f"Axon set with IPv6 failed: {result.stdout}" - assert "successfully" in result.stdout.lower() or "success" in result.stdout.lower(), \ - f"Success message not found in output: {result.stdout}" + assert ( + "successfully" in result.stdout.lower() or "success" in result.stdout.lower() + ), f"Success message not found in output: {result.stdout}" @pytest.mark.parametrize("local_chain", [None], indirect=True) @@ -273,12 +277,12 @@ def test_axon_set_invalid_inputs(local_chain, wallet_setup): """ wallet_path_charlie = "//Charlie" netuid = 1 - + # Create wallet for Charlie - keypair_charlie, wallet_charlie, wallet_path_charlie, exec_command_charlie = wallet_setup( - wallet_path_charlie + keypair_charlie, wallet_charlie, wallet_path_charlie, exec_command_charlie = ( + wallet_setup(wallet_path_charlie) ) - + # Register a subnet result = exec_command_charlie( command="subnets", @@ -296,7 +300,7 @@ def test_axon_set_invalid_inputs(local_chain, wallet_setup): ], ) assert result.exit_code == 0 - + # Register neuron result = exec_command_charlie( command="subnets", @@ -316,7 +320,7 @@ def test_axon_set_invalid_inputs(local_chain, wallet_setup): ], ) assert result.exit_code == 0 - + # Test with invalid port (too high) result = exec_command_charlie( command="axon", @@ -339,11 +343,14 @@ def test_axon_set_invalid_inputs(local_chain, wallet_setup): "--no-prompt", ], ) - + # Should fail with invalid port - assert result.exit_code != 0 or "invalid port" in result.stdout.lower() or "failed" in result.stdout.lower(), \ - f"Expected error for invalid port, got: {result.stdout}" - + assert ( + result.exit_code != 0 + or "invalid port" in result.stdout.lower() + or "failed" in result.stdout.lower() + ), f"Expected error for invalid port, got: {result.stdout}" + # Test with invalid IP result = exec_command_charlie( command="axon", @@ -366,7 +373,10 @@ def test_axon_set_invalid_inputs(local_chain, wallet_setup): "--no-prompt", ], ) - + # Should fail with invalid IP - assert result.exit_code != 0 or "invalid ip" in result.stdout.lower() or "failed" in result.stdout.lower(), \ - f"Expected error for invalid IP, got: {result.stdout}" + assert ( + result.exit_code != 0 + or "invalid ip" in result.stdout.lower() + or "failed" in result.stdout.lower() + ), f"Expected error for invalid IP, got: {result.stdout}" diff --git a/tests/unit_tests/test_axon_commands.py b/tests/unit_tests/test_axon_commands.py index 4ac4e8ec8..4ed73f9e9 100644 --- a/tests/unit_tests/test_axon_commands.py +++ b/tests/unit_tests/test_axon_commands.py @@ -1,6 +1,7 @@ """ Unit tests for axon commands (reset and set). """ + import pytest from unittest.mock import AsyncMock, MagicMock, Mock, patch from bittensor_wallet import Wallet @@ -27,7 +28,7 @@ def test_ipv6_conversion(self): # IPv6 loopback result = ip_to_int("::1") assert result == 1 - + # IPv6 address result = ip_to_int("2001:db8::1") assert result > 0 @@ -36,7 +37,7 @@ def test_invalid_ip_raises_error(self): """Test that invalid IP addresses raise errors.""" with pytest.raises(Exception): ip_to_int("invalid.ip.address") - + with pytest.raises(Exception): ip_to_int("256.256.256.256") @@ -50,24 +51,37 @@ async def test_reset_axon_success(self): # Setup mocks mock_subtensor = MagicMock() mock_subtensor.substrate.compose_call = AsyncMock(return_value="mock_call") - mock_subtensor.substrate.create_signed_extrinsic = AsyncMock(return_value="mock_extrinsic") + mock_subtensor.substrate.create_signed_extrinsic = AsyncMock( + return_value="mock_extrinsic" + ) mock_response = MagicMock() + # is_success is a property that returns a coroutine async def mock_is_success(): return True + mock_response.is_success = mock_is_success() mock_response.get_extrinsic_identifier = AsyncMock(return_value="0x123") - mock_subtensor.substrate.submit_extrinsic = AsyncMock(return_value=mock_response) - + mock_subtensor.substrate.submit_extrinsic = AsyncMock( + return_value=mock_response + ) + mock_wallet = MagicMock(spec=Wallet) - mock_wallet.hotkey.ss58_address = "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" - + mock_wallet.hotkey.ss58_address = ( + "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + ) + with ( - patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock, - patch("bittensor_cli.src.bittensor.extrinsics.serving.print_extrinsic_id", new_callable=AsyncMock), + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.unlock_key" + ) as mock_unlock, + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.print_extrinsic_id", + new_callable=AsyncMock, + ), ): mock_unlock.return_value = MagicMock(success=True) - + # Execute success, message, ext_id = await reset_axon_extrinsic( subtensor=mock_subtensor, @@ -77,12 +91,12 @@ async def mock_is_success(): wait_for_inclusion=True, wait_for_finalization=False, ) - + # Verify assert success is True assert "successfully" in message.lower() assert ext_id == "0x123" - + # Verify compose_call was called with correct parameters mock_subtensor.substrate.compose_call.assert_called_once() call_args = mock_subtensor.substrate.compose_call.call_args @@ -98,17 +112,21 @@ async def test_reset_axon_unlock_failure(self): """Test axon reset when hotkey unlock fails.""" mock_subtensor = MagicMock() mock_wallet = MagicMock(spec=Wallet) - - with patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock: - mock_unlock.return_value = MagicMock(success=False, message="Failed to unlock hotkey") - + + with patch( + "bittensor_cli.src.bittensor.extrinsics.serving.unlock_key" + ) as mock_unlock: + mock_unlock.return_value = MagicMock( + success=False, message="Failed to unlock hotkey" + ) + success, message, ext_id = await reset_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, netuid=1, prompt=False, ) - + assert success is False assert "unlock" in message.lower() assert ext_id is None @@ -118,22 +136,28 @@ async def test_reset_axon_user_cancellation(self): """Test axon reset when user cancels prompt.""" mock_subtensor = MagicMock() mock_wallet = MagicMock(spec=Wallet) - mock_wallet.hotkey.ss58_address = "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" - + mock_wallet.hotkey.ss58_address = ( + "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + ) + with ( - patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock, - patch("bittensor_cli.src.bittensor.extrinsics.serving.Confirm") as mock_confirm, + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.unlock_key" + ) as mock_unlock, + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.Confirm" + ) as mock_confirm, ): mock_unlock.return_value = MagicMock(success=True) mock_confirm.ask.return_value = False - + success, message, ext_id = await reset_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, netuid=1, prompt=True, ) - + assert success is False assert "cancelled" in message.lower() assert ext_id is None @@ -143,27 +167,37 @@ async def test_reset_axon_extrinsic_failure(self): """Test axon reset when extrinsic submission fails.""" mock_subtensor = MagicMock() mock_subtensor.substrate.compose_call = AsyncMock(return_value="mock_call") - mock_subtensor.substrate.create_signed_extrinsic = AsyncMock(return_value="mock_extrinsic") + mock_subtensor.substrate.create_signed_extrinsic = AsyncMock( + return_value="mock_extrinsic" + ) mock_response = MagicMock() + async def mock_is_success(): return False + mock_response.is_success = mock_is_success() mock_response.error_message = AsyncMock(return_value="Network error") - mock_subtensor.substrate.submit_extrinsic = AsyncMock(return_value=mock_response) - + mock_subtensor.substrate.submit_extrinsic = AsyncMock( + return_value=mock_response + ) + mock_wallet = MagicMock(spec=Wallet) - mock_wallet.hotkey.ss58_address = "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" - - with patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock: + mock_wallet.hotkey.ss58_address = ( + "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + ) + + with patch( + "bittensor_cli.src.bittensor.extrinsics.serving.unlock_key" + ) as mock_unlock: mock_unlock.return_value = MagicMock(success=True) - + success, message, ext_id = await reset_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, netuid=1, prompt=False, ) - + assert success is False assert len(message) > 0 assert ext_id is None @@ -177,23 +211,36 @@ async def test_set_axon_success(self): """Test successful axon set.""" mock_subtensor = MagicMock() mock_subtensor.substrate.compose_call = AsyncMock(return_value="mock_call") - mock_subtensor.substrate.create_signed_extrinsic = AsyncMock(return_value="mock_extrinsic") + mock_subtensor.substrate.create_signed_extrinsic = AsyncMock( + return_value="mock_extrinsic" + ) mock_response = MagicMock() + async def mock_is_success(): return True + mock_response.is_success = mock_is_success() mock_response.get_extrinsic_identifier = AsyncMock(return_value="0x123") - mock_subtensor.substrate.submit_extrinsic = AsyncMock(return_value=mock_response) - + mock_subtensor.substrate.submit_extrinsic = AsyncMock( + return_value=mock_response + ) + mock_wallet = MagicMock(spec=Wallet) - mock_wallet.hotkey.ss58_address = "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" - + mock_wallet.hotkey.ss58_address = ( + "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + ) + with ( - patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock, - patch("bittensor_cli.src.bittensor.extrinsics.serving.print_extrinsic_id", new_callable=AsyncMock), + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.unlock_key" + ) as mock_unlock, + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.print_extrinsic_id", + new_callable=AsyncMock, + ), ): mock_unlock.return_value = MagicMock(success=True) - + success, message, ext_id = await set_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, @@ -206,12 +253,12 @@ async def mock_is_success(): wait_for_inclusion=True, wait_for_finalization=False, ) - + assert success is True assert "successfully" in message.lower() assert ext_id == "0x123" assert "192.168.1.100:8091" in message - + # Verify compose_call was called with correct parameters mock_subtensor.substrate.compose_call.assert_called_once() call_args = mock_subtensor.substrate.compose_call.call_args @@ -227,7 +274,7 @@ async def test_set_axon_invalid_port(self): """Test axon set with invalid port number.""" mock_subtensor = MagicMock() mock_wallet = MagicMock(spec=Wallet) - + # Test port too high success, message, ext_id = await set_axon_extrinsic( subtensor=mock_subtensor, @@ -237,11 +284,11 @@ async def test_set_axon_invalid_port(self): port=70000, prompt=False, ) - + assert success is False assert "Invalid port" in message assert ext_id is None - + # Test negative port success, message, ext_id = await set_axon_extrinsic( subtensor=mock_subtensor, @@ -251,7 +298,7 @@ async def test_set_axon_invalid_port(self): port=-1, prompt=False, ) - + assert success is False assert "Invalid port" in message assert ext_id is None @@ -261,7 +308,7 @@ async def test_set_axon_invalid_ip(self): """Test axon set with invalid IP address.""" mock_subtensor = MagicMock() mock_wallet = MagicMock(spec=Wallet) - + success, message, ext_id = await set_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, @@ -270,7 +317,7 @@ async def test_set_axon_invalid_ip(self): port=8091, prompt=False, ) - + assert success is False assert "Invalid IP" in message assert ext_id is None @@ -280,10 +327,14 @@ async def test_set_axon_unlock_failure(self): """Test axon set when hotkey unlock fails.""" mock_subtensor = MagicMock() mock_wallet = MagicMock(spec=Wallet) - - with patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock: - mock_unlock.return_value = MagicMock(success=False, message="Failed to unlock hotkey") - + + with patch( + "bittensor_cli.src.bittensor.extrinsics.serving.unlock_key" + ) as mock_unlock: + mock_unlock.return_value = MagicMock( + success=False, message="Failed to unlock hotkey" + ) + success, message, ext_id = await set_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, @@ -292,7 +343,7 @@ async def test_set_axon_unlock_failure(self): port=8091, prompt=False, ) - + assert success is False assert "unlock" in message.lower() assert "Failed to unlock hotkey" in message @@ -302,15 +353,21 @@ async def test_set_axon_user_cancellation(self): """Test axon set when user cancels prompt.""" mock_subtensor = MagicMock() mock_wallet = MagicMock(spec=Wallet) - mock_wallet.hotkey.ss58_address = "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" - + mock_wallet.hotkey.ss58_address = ( + "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + ) + with ( - patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock, - patch("bittensor_cli.src.bittensor.extrinsics.serving.Confirm") as mock_confirm, + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.unlock_key" + ) as mock_unlock, + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.Confirm" + ) as mock_confirm, ): mock_unlock.return_value = MagicMock(success=True) mock_confirm.ask.return_value = False - + success, message, ext_id = await set_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, @@ -319,7 +376,7 @@ async def test_set_axon_user_cancellation(self): port=8091, prompt=True, ) - + assert success is False assert "cancelled" in message.lower() @@ -328,23 +385,36 @@ async def test_set_axon_with_ipv6(self): """Test axon set with IPv6 address.""" mock_subtensor = MagicMock() mock_subtensor.substrate.compose_call = AsyncMock(return_value="mock_call") - mock_subtensor.substrate.create_signed_extrinsic = AsyncMock(return_value="mock_extrinsic") + mock_subtensor.substrate.create_signed_extrinsic = AsyncMock( + return_value="mock_extrinsic" + ) mock_response = MagicMock() + async def mock_is_success(): return True + mock_response.is_success = mock_is_success() mock_response.get_extrinsic_identifier = AsyncMock(return_value="0x123") - mock_subtensor.substrate.submit_extrinsic = AsyncMock(return_value=mock_response) - + mock_subtensor.substrate.submit_extrinsic = AsyncMock( + return_value=mock_response + ) + mock_wallet = MagicMock(spec=Wallet) - mock_wallet.hotkey.ss58_address = "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" - + mock_wallet.hotkey.ss58_address = ( + "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + ) + with ( - patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock, - patch("bittensor_cli.src.bittensor.extrinsics.serving.print_extrinsic_id", new_callable=AsyncMock), + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.unlock_key" + ) as mock_unlock, + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.print_extrinsic_id", + new_callable=AsyncMock, + ), ): mock_unlock.return_value = MagicMock(success=True) - + success, message, ext_id = await set_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, @@ -355,7 +425,7 @@ async def mock_is_success(): protocol=4, prompt=False, ) - + assert success is True assert "successfully" in message.lower() assert ext_id == "0x123" @@ -370,13 +440,17 @@ async def test_set_axon_exception_handling(self): mock_subtensor.substrate.compose_call = AsyncMock( side_effect=Exception("Unexpected error") ) - + mock_wallet = MagicMock(spec=Wallet) - mock_wallet.hotkey.ss58_address = "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" - - with patch("bittensor_cli.src.bittensor.extrinsics.serving.unlock_key") as mock_unlock: + mock_wallet.hotkey.ss58_address = ( + "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + ) + + with patch( + "bittensor_cli.src.bittensor.extrinsics.serving.unlock_key" + ) as mock_unlock: mock_unlock.return_value = MagicMock(success=True) - + success, message, ext_id = await set_axon_extrinsic( subtensor=mock_subtensor, wallet=mock_wallet, @@ -385,7 +459,7 @@ async def test_set_axon_exception_handling(self): port=8091, prompt=False, ) - + assert success is False assert len(message) > 0 assert ext_id is None @@ -398,12 +472,10 @@ class TestAxonCLICommands: def test_axon_reset_command_handler(self, mock_axon): """Test axon reset CLI command handler.""" from bittensor_cli.cli import CLIManager - + cli_manager = CLIManager() - mock_axon.reset = AsyncMock( - return_value=None - ) - + mock_axon.reset = AsyncMock(return_value=None) + with ( patch.object(cli_manager, "verbosity_handler"), patch.object(cli_manager, "wallet_ask") as mock_wallet_ask, @@ -414,7 +486,7 @@ def test_axon_reset_command_handler(self, mock_axon): mock_wallet_ask.return_value = mock_wallet mock_subtensor = Mock() mock_init_chain.return_value = mock_subtensor - + cli_manager.axon_reset( netuid=1, wallet_name="test_wallet", @@ -428,10 +500,10 @@ def test_axon_reset_command_handler(self, mock_axon): verbose=False, json_output=False, ) - + # Verify wallet_ask was called correctly mock_wallet_ask.assert_called_once() - + # Verify _run_command was called mock_run_command.assert_called_once() @@ -439,12 +511,10 @@ def test_axon_reset_command_handler(self, mock_axon): def test_axon_set_command_handler(self, mock_axon): """Test axon set CLI command handler.""" from bittensor_cli.cli import CLIManager - + cli_manager = CLIManager() - mock_axon.set_axon = AsyncMock( - return_value=None - ) - + mock_axon.set_axon = AsyncMock(return_value=None) + with ( patch.object(cli_manager, "verbosity_handler"), patch.object(cli_manager, "wallet_ask") as mock_wallet_ask, @@ -455,7 +525,7 @@ def test_axon_set_command_handler(self, mock_axon): mock_wallet_ask.return_value = mock_wallet mock_subtensor = Mock() mock_init_chain.return_value = mock_subtensor - + cli_manager.axon_set( netuid=1, ip="192.168.1.100", @@ -473,9 +543,9 @@ def test_axon_set_command_handler(self, mock_axon): verbose=False, json_output=False, ) - + # Verify wallet_ask was called correctly mock_wallet_ask.assert_called_once() - + # Verify _run_command was called mock_run_command.assert_called_once() From 00d0d501553e8be1e375ea3d2db3bc4691fcc03f Mon Sep 17 00:00:00 2001 From: Mobile-Crest Date: Fri, 5 Dec 2025 03:56:12 +0100 Subject: [PATCH 5/6] fix e2e test --- tests/e2e_tests/test_axon.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/e2e_tests/test_axon.py b/tests/e2e_tests/test_axon.py index 5b3bf539f..e2a598cc7 100644 --- a/tests/e2e_tests/test_axon.py +++ b/tests/e2e_tests/test_axon.py @@ -44,6 +44,8 @@ def test_axon_reset_and_set(local_chain, wallet_setup): wallet_alice.name, "--wallet-hotkey", wallet_alice.hotkey_str, + "--subnet-name", + "Test Axon Subnet", "--no-prompt", ], ) @@ -211,6 +213,8 @@ def test_axon_set_with_ipv6(local_chain, wallet_setup): wallet_bob.name, "--wallet-hotkey", wallet_bob.hotkey_str, + "--subnet-name", + "Test IPv6 Subnet", "--no-prompt", ], ) @@ -296,6 +300,8 @@ def test_axon_set_invalid_inputs(local_chain, wallet_setup): wallet_charlie.name, "--wallet-hotkey", wallet_charlie.hotkey_str, + "--subnet-name", + "Test Invalid Inputs Subnet", "--no-prompt", ], ) From e4b3d32491f472c1b2c623948284dd132a4036c8 Mon Sep 17 00:00:00 2001 From: Mobile-Crest Date: Fri, 5 Dec 2025 09:27:27 +0100 Subject: [PATCH 6/6] fix e2e test error --- tests/e2e_tests/test_axon.py | 42 ++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/tests/e2e_tests/test_axon.py b/tests/e2e_tests/test_axon.py index e2a598cc7..bf070e7c7 100644 --- a/tests/e2e_tests/test_axon.py +++ b/tests/e2e_tests/test_axon.py @@ -46,6 +46,20 @@ def test_axon_reset_and_set(local_chain, wallet_setup): wallet_alice.hotkey_str, "--subnet-name", "Test Axon Subnet", + "--repo", + "https://github.com/test/axon-subnet", + "--contact", + "test@opentensor.dev", + "--url", + "https://testaxon.com", + "--discord", + "test#1234", + "--description", + "Test subnet for axon e2e testing", + "--logo-url", + "https://testaxon.com/logo.png", + "--additional-info", + "Axon test subnet", "--no-prompt", ], ) @@ -215,6 +229,20 @@ def test_axon_set_with_ipv6(local_chain, wallet_setup): wallet_bob.hotkey_str, "--subnet-name", "Test IPv6 Subnet", + "--repo", + "https://github.com/test/ipv6-subnet", + "--contact", + "ipv6@opentensor.dev", + "--url", + "https://testipv6.com", + "--discord", + "ipv6#5678", + "--description", + "Test subnet for IPv6 axon testing", + "--logo-url", + "https://testipv6.com/logo.png", + "--additional-info", + "IPv6 test subnet", "--no-prompt", ], ) @@ -302,6 +330,20 @@ def test_axon_set_invalid_inputs(local_chain, wallet_setup): wallet_charlie.hotkey_str, "--subnet-name", "Test Invalid Inputs Subnet", + "--repo", + "https://github.com/test/invalid-subnet", + "--contact", + "invalid@opentensor.dev", + "--url", + "https://testinvalid.com", + "--discord", + "invalid#9999", + "--description", + "Test subnet for invalid inputs testing", + "--logo-url", + "https://testinvalid.com/logo.png", + "--additional-info", + "Invalid inputs test subnet", "--no-prompt", ], )