Merge pull request #2766 from opentensor/feat/roman/add-get-owned-hotkeys

Add get_owned_hotkeys to subtensor and async one + tests

Author: basfroman

bittensor/core/async_subtensor.py

@@ -1572,6 +1572,36 @@ async def get_neuron_for_pubkey_and_subnet(
                 reuse_block=reuse_block,
             )
 
+    async def get_owned_hotkeys(
+        self,
+        coldkey_ss58: str,
+        block: Optional[int] = None,
+        block_hash: Optional[str] = None,
+        reuse_block: bool = False,
+    ) -> list[str]:
+        """
+        Retrieves all hotkeys owned by a specific coldkey address.
+
+        Args:
+            coldkey_ss58 (str): The SS58 address of the coldkey to query.
+            block (int): The blockchain block number for the query.
+            block_hash (str): The hash of the blockchain block number for the query.
+            reuse_block (bool): Whether to reuse the last-used blockchain block hash.
+
+        Returns:
+            list[str]: A list of hotkey SS58 addresses owned by the coldkey.
+        """
+        block_hash = await self.determine_block_hash(block, block_hash, reuse_block)
+        owned_hotkeys = await self.substrate.query(
+            module="SubtensorModule",
+            storage_function="OwnedHotkeys",
+            params=[coldkey_ss58],
+            block_hash=block_hash,
+            reuse_block_hash=reuse_block,
+        )
+
+        return [decode_account_id(hotkey[0]) for hotkey in owned_hotkeys or []]
+
     async def get_stake(
         self,
         coldkey_ss58: str,

bittensor/core/subtensor.py

@@ -1200,6 +1200,33 @@ def get_neuron_for_pubkey_and_subnet(
 
         return NeuronInfo.from_dict(result)
 
+    def get_owned_hotkeys(
+        self,
+        coldkey_ss58: str,
+        block: Optional[int] = None,
+        reuse_block: bool = False,
+    ) -> list[str]:
+        """
+        Retrieves all hotkeys owned by a specific coldkey address.
+
+        Args:
+            coldkey_ss58 (str): The SS58 address of the coldkey to query.
+            block (int): The blockchain block number for the query.
+            reuse_block (bool): Whether to reuse the last-used blockchain block hash.
+
+        Returns:
+            list[str]: A list of hotkey SS58 addresses owned by the coldkey.
+        """
+        block_hash = self.determine_block_hash(block)
+        owned_hotkeys = self.substrate.query(
+            module="SubtensorModule",
+            storage_function="OwnedHotkeys",
+            params=[coldkey_ss58],
+            block_hash=block_hash,
+            reuse_block_hash=reuse_block,
+        )
+        return [decode_account_id(hotkey[0]) for hotkey in owned_hotkeys or []]
+
     def get_stake(
         self,
         coldkey_ss58: str,

tests/unit_tests/test_async_subtensor.py

@@ -2956,3 +2956,59 @@ async def test_get_timestamp(mocker, subtensor):
     )
     actual_result = await subtensor.get_timestamp(block=fake_block)
     assert expected_result == actual_result
+
+
+@pytest.mark.asyncio
+async def test_get_owned_hotkeys_happy_path(subtensor, mocker):
+    """Tests that the output of get_owned_hotkeys."""
+    # Prep
+    fake_coldkey = "fake_hotkey"
+    fake_hotkey = "fake_hotkey"
+    fake_hotkeys = [
+        [
+            fake_hotkey,
+        ]
+    ]
+    mocked_subtensor = mocker.AsyncMock(return_value=fake_hotkeys)
+    mocker.patch.object(subtensor.substrate, "query", new=mocked_subtensor)
+
+    mocked_decode_account_id = mocker.Mock()
+    mocker.patch.object(
+        async_subtensor, "decode_account_id", new=mocked_decode_account_id
+    )
+
+    # Call
+    result = await subtensor.get_owned_hotkeys(fake_coldkey)
+
+    # Asserts
+    mocked_subtensor.assert_awaited_once_with(
+        module="SubtensorModule",
+        storage_function="OwnedHotkeys",
+        params=[fake_coldkey],
+        block_hash=None,
+        reuse_block_hash=False,
+    )
+    assert result == [mocked_decode_account_id.return_value]
+    mocked_decode_account_id.assert_called_once_with(fake_hotkey)
+
+
+@pytest.mark.asyncio
+async def test_get_owned_hotkeys_return_empty(subtensor, mocker):
+    """Tests that the output of get_owned_hotkeys is empty."""
+    # Prep
+    fake_coldkey = "fake_hotkey"
+    mocked_subtensor = mocker.AsyncMock(return_value=[])
+    mocker.patch.object(subtensor.substrate, "query", new=mocked_subtensor)
+
+    # Call
+    result = await subtensor.get_owned_hotkeys(fake_coldkey)
+
+    # Asserts
+    mocked_subtensor.assert_awaited_once_with(
+        module="SubtensorModule",
+        storage_function="OwnedHotkeys",
+        params=[fake_coldkey],
+        block_hash=None,
+        reuse_block_hash=False,
+    )
+    assert result == []

tests/unit_tests/test_subtensor.py

@@ -3332,3 +3332,57 @@ def test_stake_fee_methods(mocker, subtensor):
         ],
         block=None,
     )
+
+
+def test_get_owned_hotkeys_happy_path(subtensor, mocker):
+    """Tests that the output of get_owned_hotkeys."""
+    # Prep
+    fake_coldkey = "fake_hotkey"
+    fake_hotkey = "fake_hotkey"
+    fake_hotkeys = [
+        [
+            fake_hotkey,
+        ]
+    ]
+    mocked_subtensor = mocker.Mock(return_value=fake_hotkeys)
+    mocker.patch.object(subtensor.substrate, "query", new=mocked_subtensor)
+
+    mocked_decode_account_id = mocker.Mock()
+    mocker.patch.object(
+        subtensor_module, "decode_account_id", new=mocked_decode_account_id
+    )
+
+    # Call
+    result = subtensor.get_owned_hotkeys(fake_coldkey)
+
+    # Asserts
+    mocked_subtensor.assert_called_once_with(
+        module="SubtensorModule",
+        storage_function="OwnedHotkeys",
+        params=[fake_coldkey],
+        block_hash=None,
+        reuse_block_hash=False,
+    )
+    assert result == [mocked_decode_account_id.return_value]
+    mocked_decode_account_id.assert_called_once_with(fake_hotkey)
+
+
+def test_get_owned_hotkeys_return_empty(subtensor, mocker):
+    """Tests that the output of get_owned_hotkeys is empty."""
+    # Prep
+    fake_coldkey = "fake_hotkey"
+    mocked_subtensor = mocker.Mock(return_value=[])
+    mocker.patch.object(subtensor.substrate, "query", new=mocked_subtensor)
+
+    # Call
+    result = subtensor.get_owned_hotkeys(fake_coldkey)
+
+    # Asserts
+    mocked_subtensor.assert_called_once_with(
+        module="SubtensorModule",
+        storage_function="OwnedHotkeys",
+        params=[fake_coldkey],
+        block_hash=None,
+        reuse_block_hash=False,
+    )
+    assert result == []