Merge pull request #401 from opentensor/feat/associate-hotkey

ibraheem-abe · web-flow · commit aad75214b011 · 2025-03-21T14:45:52.000-07:00
Feat/associate hotkey
diff --git a/bittensor_cli/cli.py b/bittensor_cli/cli.py
@@ -716,6 +716,9 @@ def __init__(self):
         self.wallet_app.command(
             "new-coldkey", rich_help_panel=HELP_PANELS["WALLET"]["MANAGEMENT"]
         )(self.wallet_new_coldkey)
+        self.wallet_app.command(
+            "associate-hotkey", rich_help_panel=HELP_PANELS["WALLET"]["MANAGEMENT"]
+        )(self.wallet_associate_hotkey)
         self.wallet_app.command(
             "create", rich_help_panel=HELP_PANELS["WALLET"]["MANAGEMENT"]
         )(self.wallet_create_wallet)
@@ -2265,6 +2268,74 @@ def wallet_new_hotkey(
             wallets.new_hotkey(wallet, n_words, use_password, uri, overwrite)
         )
 
+    def wallet_associate_hotkey(
+        self,
+        wallet_name: Optional[str] = Options.wallet_name,
+        wallet_path: Optional[str] = Options.wallet_path,
+        wallet_hotkey: Optional[str] = Options.wallet_hotkey_ss58,
+        network: Optional[list[str]] = Options.network,
+        prompt: bool = Options.prompt,
+        quiet: bool = Options.quiet,
+        verbose: bool = Options.verbose,
+    ):
+        """
+        Associate a hotkey with a wallet(coldkey).
+
+        USAGE
+
+        This command is used to associate a hotkey with a wallet(coldkey).
+
+        EXAMPLE
+
+        [green]$[/green] btcli wallet associate-hotkey --hotkey-name hotkey_name
+        [green]$[/green] btcli wallet associate-hotkey --hotkey-ss58 5DkQ4...
+        """
+        self.verbosity_handler(quiet, verbose)
+        if not wallet_name:
+            wallet_name = Prompt.ask(
+                "Enter the [blue]wallet name[/blue] [dim](which you want to associate with the hotkey)[/dim]",
+                default=self.config.get("wallet_name") or defaults.wallet.name,
+            )
+        if not wallet_hotkey:
+            wallet_hotkey = Prompt.ask(
+                "Enter the [blue]hotkey[/blue] name or "
+                "[blue]hotkey ss58 address[/blue] [dim](to associate with your coldkey)[/dim]"
+            )
+
+        hotkey_display = None
+        if is_valid_ss58_address(wallet_hotkey):
+            hotkey_ss58 = wallet_hotkey
+            wallet = self.wallet_ask(
+                wallet_name,
+                wallet_path,
+                None,
+                ask_for=[WO.NAME, WO.PATH],
+                validate=WV.WALLET,
+            )
+            hotkey_display = (
+                f"hotkey [{COLORS.GENERAL.HK}]{hotkey_ss58}[/{COLORS.GENERAL.HK}]"
+            )
+        else:
+            wallet = self.wallet_ask(
+                wallet_name,
+                wallet_path,
+                wallet_hotkey,
+                ask_for=[WO.NAME, WO.PATH, WO.HOTKEY],
+                validate=WV.WALLET_AND_HOTKEY,
+            )
+            hotkey_ss58 = wallet.hotkey.ss58_address
+            hotkey_display = f"hotkey [blue]{wallet_hotkey}[/blue] [{COLORS.GENERAL.HK}]({hotkey_ss58})[/{COLORS.GENERAL.HK}]"
+
+        return self._run_command(
+            wallets.associate_hotkey(
+                wallet,
+                self.initialize_chain(network),
+                hotkey_ss58,
+                hotkey_display,
+                prompt,
+            )
+        )
+
     def wallet_new_coldkey(
         self,
         wallet_name: Optional[str] = Options.wallet_name,
diff --git a/bittensor_cli/src/commands/wallets.py b/bittensor_cli/src/commands/wallets.py
@@ -50,6 +50,73 @@
 )
 
 
+async def associate_hotkey(
+    wallet: Wallet,
+    subtensor: SubtensorInterface,
+    hotkey_ss58: str,
+    hotkey_display: str,
+    prompt: bool = False,
+):
+    """Associates a hotkey with a wallet"""
+
+    owner_ss58 = await subtensor.get_hotkey_owner(hotkey_ss58)
+    if owner_ss58:
+        if owner_ss58 == wallet.coldkeypub.ss58_address:
+            console.print(
+                f":white_heavy_check_mark: {hotkey_display.capitalize()} is already "
+                f"associated with \nwallet [blue]{wallet.name}[/blue], "
+                f"SS58: [{COLORS.GENERAL.CK}]{owner_ss58}[/{COLORS.GENERAL.CK}]"
+            )
+            return True
+        else:
+            owner_wallet = _get_wallet_by_ss58(wallet.path, owner_ss58)
+            wallet_name = owner_wallet.name if owner_wallet else "unknown wallet"
+            console.print(
+                f"[yellow]Warning[/yellow]: {hotkey_display.capitalize()} is already associated with \n"
+                f"wallet: [blue]{wallet_name}[/blue], SS58: [{COLORS.GENERAL.CK}]{owner_ss58}[/{COLORS.GENERAL.CK}]"
+            )
+            return False
+    else:
+        console.print(
+            f"{hotkey_display.capitalize()} is not associated with any wallet"
+        )
+
+    if prompt and not Confirm.ask("Do you want to continue with the association?"):
+        return False
+
+    if not unlock_key(wallet).success:
+        return False
+
+    call = await subtensor.substrate.compose_call(
+        call_module="SubtensorModule",
+        call_function="try_associate_hotkey",
+        call_params={
+            "hotkey": hotkey_ss58,
+        },
+    )
+
+    with console.status(":satellite: Associating hotkey on-chain..."):
+        success, err_msg = await subtensor.sign_and_send_extrinsic(
+            call,
+            wallet,
+            wait_for_inclusion=True,
+            wait_for_finalization=False,
+        )
+
+        if not success:
+            console.print(
+                f"[red]:cross_mark: Failed to associate hotkey: {err_msg}[/red]"
+            )
+            return False
+
+        console.print(
+            f":white_heavy_check_mark: Successfully associated {hotkey_display} with \n"
+            f"wallet [blue]{wallet.name}[/blue], "
+            f"SS58: [{COLORS.GENERAL.CK}]{wallet.coldkeypub.ss58_address}[/{COLORS.GENERAL.CK}]"
+        )
+        return True
+
+
 async def regen_coldkey(
     wallet: Wallet,
     mnemonic: Optional[str],
@@ -257,6 +324,15 @@ def get_coldkey_wallets_for_path(path: str) -> list[Wallet]:
     return wallets
 
 
+def _get_wallet_by_ss58(path: str, ss58_address: str) -> Optional[Wallet]:
+    """Find a wallet by its SS58 address in the given path."""
+    ss58_addresses, wallet_names = _get_coldkey_ss58_addresses_for_path(path)
+    for wallet_name, addr in zip(wallet_names, ss58_addresses):
+        if addr == ss58_address:
+            return Wallet(path=path, name=wallet_name)
+    return None
+
+
 def _get_coldkey_ss58_addresses_for_path(path: str) -> tuple[list[str], list[str]]:
     """Get all coldkey ss58 addresses from path."""
 
@@ -1596,8 +1672,7 @@ async def find_coldkey_swap_extrinsic(
     """
 
     current_block, genesis_block = await asyncio.gather(
-        subtensor.substrate.get_block_number(),
-        subtensor.substrate.get_block_hash(0)
+        subtensor.substrate.get_block_number(), subtensor.substrate.get_block_hash(0)
     )
     if (
         current_block - start_block > 300
diff --git a/tests/e2e_tests/test_wallet_interactions.py b/tests/e2e_tests/test_wallet_interactions.py
@@ -510,3 +510,112 @@ def test_wallet_identities(local_chain, wallet_setup):
     assert "Message signed successfully" in sign_using_coldkey.stdout
 
     print("✅ Passed wallet set-id, get-id, sign command")
+
+
+def test_wallet_associate_hotkey(local_chain, wallet_setup):
+    """
+    Test the associating hotkeys and their different cases.
+
+    Steps:
+        1. Create wallets for Alice, Bob, and Charlie
+        2. Associate a hotkey with Alice's wallet using hotkey name
+        3. Verify the association is successful
+        4. Try to associate Alice's hotkey with Bob's wallet (should fail)
+        5. Try to associate Alice's hotkey again (should show already associated)
+        6. Associate Charlie's hotkey with Bob's wallet using SS58 address
+
+    Raises:
+        AssertionError: If any of the checks or verifications fail
+    """
+    print("Testing wallet associate-hotkey command 🧪")
+
+    _, wallet_alice, wallet_path_alice, exec_command_alice = wallet_setup("//Alice")
+    _, wallet_bob, wallet_path_bob, exec_command_bob = wallet_setup("//Bob")
+    _, wallet_charlie, _, _ = wallet_setup("//Charlie")
+
+    # Associate Alice's default hotkey with her wallet
+    result = exec_command_alice(
+        command="wallet",
+        sub_command="associate-hotkey",
+        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 successful association
+    assert "Successfully associated hotkey" in result.stdout
+    assert wallet_alice.hotkey.ss58_address in result.stdout
+    assert wallet_alice.coldkeypub.ss58_address in result.stdout
+    assert wallet_alice.hotkey_str in result.stdout
+
+    # Try to associate Alice's hotkey with Bob's wallet (should fail)
+    result = exec_command_bob(
+        command="wallet",
+        sub_command="associate-hotkey",
+        extra_args=[
+            "--wallet-path",
+            wallet_path_bob,
+            "--chain",
+            "ws://127.0.0.1:9945",
+            "--wallet-name",
+            wallet_bob.name,
+            "--hotkey-ss58",
+            wallet_alice.hotkey.ss58_address,
+            "--no-prompt",
+        ],
+    )
+
+    assert "Warning" in result.stdout
+    assert "is already associated with" in result.stdout
+
+    # Try to associate Alice's hotkey again with Alice's wallet
+    result = exec_command_alice(
+        command="wallet",
+        sub_command="associate-hotkey",
+        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 "is already associated with" in result.stdout
+    assert "wallet" in result.stdout
+    assert wallet_alice.name in result.stdout
+
+    # Associate Charlie's hotkey with Bob's wallet using SS58 address
+    result = exec_command_bob(
+        command="wallet",
+        sub_command="associate-hotkey",
+        extra_args=[
+            "--wallet-path",
+            wallet_path_bob,
+            "--chain",
+            "ws://127.0.0.1:9945",
+            "--wallet-name",
+            wallet_bob.name,
+            "--hotkey-ss58",
+            wallet_charlie.hotkey.ss58_address,
+            "--no-prompt",
+        ],
+    )
+
+    assert "Successfully associated hotkey" in result.stdout
+    assert wallet_charlie.hotkey.ss58_address in result.stdout
+    assert wallet_bob.coldkeypub.ss58_address in result.stdout
+
+    print("✅ Passed wallet associate-hotkey command")
