fix test + ruff

Author: basfroman

bittensor/core/chain_data/dynamic_info.py

@@ -68,7 +68,7 @@ def fix_decoded_values(cls, decoded: dict) -> "DynamicInfo":
         netuid = int(decoded["netuid"])
         symbol = bytes([int(b) for b in decoded["token_symbol"]]).decode()
         subnet_name = bytes([int(b) for b in decoded["subnet_name"]]).decode()
-        
+
         is_dynamic = (
             True if int(decoded["netuid"]) > 0 else False
         )  # Root is not dynamic

bittensor/core/chain_data/stake_info.py

@@ -32,7 +32,6 @@ class StakeInfo:
     drain: int
     is_registered: bool
 
-
     @classmethod
     def fix_decoded_values(cls, decoded: Any) -> "StakeInfo":
         """Fixes the decoded values."""

bittensor/core/extrinsics/unstaking.py

@@ -244,6 +244,7 @@ def unstake_multiple_extrinsic(
         subtensor (bittensor.core.subtensor.Subtensor): Subtensor instance.
         wallet (bittensor_wallet.Wallet): The wallet with the coldkey to unstake to.
         hotkey_ss58s (List[str]): List of hotkeys to unstake from.
+        netuids (List[int]): List of netuids to unstake from.
         amounts (List[Union[Balance, float]]): List of amounts to unstake. If ``None``, unstake all.
         wait_for_inclusion (bool): If set, waits for the extrinsic to enter a block before returning ``true``, or returns ``false`` if the extrinsic fails to enter the block within the timeout.
         wait_for_finalization (bool): If set, waits for the extrinsic to be finalized on the chain before returning ``true``, or returns ``false`` if the extrinsic fails to be finalized within the timeout.

bittensor/core/subtensor.py

@@ -1703,12 +1703,15 @@ def get_delegate_by_hotkey(
 
         return DelegateInfo.from_vec_u8(bytes(result))
 
-    def get_stake_for_coldkey(self, coldkey_ss58: str) -> Optional[list["StakeInfo"]]:
+    def get_stake_for_coldkey(
+        self, coldkey_ss58: str, block: Optional[int] = None
+    ) -> Optional[list["StakeInfo"]]:
         """
         Retrieves the stake information for a given coldkey.
 
         Args:
             coldkey_ss58 (str): The SS58 address of the coldkey.
+            block (Optional[int]): The block number at which to query the stake information.
 
         Returns:
             Optional[list[StakeInfo]]: A list of StakeInfo objects, or ``None`` if no stake information is found.
@@ -1718,6 +1721,7 @@ def get_stake_for_coldkey(self, coldkey_ss58: str) -> Optional[list["StakeInfo"]
             runtime_api="StakeInfoRuntimeApi",
             method="get_stake_info_for_coldkey",
             params=[encoded_coldkey],
+            block=block,
         )
 
         if hex_bytes_result is None:
@@ -1730,7 +1734,11 @@ def get_stake_for_coldkey(self, coldkey_ss58: str) -> Optional[list["StakeInfo"]
         return StakeInfo.list_from_vec_u8(bytes_result)
 
     def get_stake_for_coldkey_and_hotkey(
-        self, hotkey_ss58: str, coldkey_ss58: str, netuid: Optional[int] = None
+        self,
+        hotkey_ss58: str,
+        coldkey_ss58: str,
+        netuid: Optional[int] = None,
+        block: Optional[int] = None,
     ) -> Optional[Union["StakeInfo", list["StakeInfo"]]]:
         """
         Returns the stake under a coldkey - hotkey pairing.
@@ -1739,11 +1747,12 @@ def get_stake_for_coldkey_and_hotkey(
             hotkey_ss58 (str): The SS58 address of the hotkey.
             coldkey_ss58 (str): The SS58 address of the coldkey.
             netuid (Optional[int]): The subnet ID to filter by. If provided, only returns stake for this specific subnet.
+            block (Optional[int]): The block number at which to query the stake information.
 
         Returns:
             Optional[StakeInfo]: The StakeInfo object/s under the coldkey - hotkey pairing, or ``None`` if the pairing does not exist or the stake is not found.
         """
-        all_stakes = self.get_stake_for_coldkey(coldkey_ss58)
+        all_stakes = self.get_stake_for_coldkey(coldkey_ss58, block)
         stakes = [
             stake
             for stake in all_stakes
@@ -2426,6 +2435,7 @@ def unstake_multiple(
         Args:
             wallet (bittensor_wallet.Wallet): The wallet linked to the coldkey from which the stakes are being withdrawn.
             hotkey_ss58s (List[str]): A list of hotkey ``SS58`` addresses to unstake from.
+            netuids (List[int]): A list of neuron netuids to unstake from.
             amounts (List[Union[Balance, float]]): The amounts of TAO to unstake from each hotkey. If not provided, unstakes all available stakes.
             wait_for_inclusion (bool): Waits for the transaction to be included in a block.
             wait_for_finalization (bool): Waits for the transaction to be finalized on the blockchain.

tests/unit_tests/test_subtensor.py

@@ -2196,45 +2196,78 @@ def test_networks_during_connection(mocker):
         sub.chain_endpoint = settings.NETWORK_MAP.get(network)
 
 
-@pytest.mark.parametrize(
-    "fake_value_result",
-    [1, None],
-    ids=["result has value attr", "result has not value attr"],
-)
-def test_get_stake_for_coldkey_and_hotkey(subtensor, mocker, fake_value_result):
-    """Test get_stake_for_coldkey_and_hotkey calls right method with correct arguments."""
+def test_get_stake_for_coldkey_and_hotkey_with_single_result(subtensor, mocker):
+    """Test `get_stake_for_coldkey_and_hotkey` calls right method with correct arguments and get 1 stake info."""
     # Preps
     fake_hotkey_ss58 = "FAKE_H_SS58"
     fake_coldkey_ss58 = "FAKE_C_SS58"
+    fake_netuid = 255
     fake_block = 123
 
-    return_value = (
-        mocker.Mock(value=fake_value_result)
-        if fake_value_result is not None
-        else fake_value_result
+    fake_stake_info_1 = mocker.Mock(hotkey_ss58="some")
+    fake_stake_info_2 = mocker.Mock(
+        hotkey_ss58=fake_hotkey_ss58, netuid=fake_netuid, stake=100
     )
 
-    subtensor.query_subtensor = mocker.patch.object(
-        subtensor, "query_subtensor", return_value=return_value
+    return_value = [
+        fake_stake_info_1,
+        fake_stake_info_2,
+    ]
+
+    subtensor.get_stake_for_coldkey = mocker.patch.object(
+        subtensor, "get_stake_for_coldkey", return_value=return_value
     )
-    spy_balance_from_rao = mocker.spy(subtensor_module.Balance, "from_rao")
 
     # Call
     result = subtensor.get_stake_for_coldkey_and_hotkey(
         hotkey_ss58=fake_hotkey_ss58,
         coldkey_ss58=fake_coldkey_ss58,
+        netuid=fake_netuid,
         block=fake_block,
     )
 
     # Asserts
-    subtensor.query_subtensor.assert_called_once_with(
-        "Stake", fake_block, [fake_hotkey_ss58, fake_coldkey_ss58]
+    subtensor.get_stake_for_coldkey.assert_called_once_with(
+        fake_coldkey_ss58, fake_block
+    )
+    assert result == fake_stake_info_2
+
+
+def test_get_stake_for_coldkey_and_hotkey_with_multiple_result(subtensor, mocker):
+    """Test `get_stake_for_coldkey_and_hotkey` calls right method with correct arguments and get multiple stake info."""
+    # Preps
+    fake_hotkey_ss58 = "FAKE_H_SS58"
+    fake_coldkey_ss58 = "FAKE_C_SS58"
+    fake_netuid = 255
+    fake_block = 123
+
+    fake_stake_info_1 = mocker.Mock(hotkey_ss58="some")
+    fake_stake_info_2 = mocker.Mock(
+        hotkey_ss58=fake_hotkey_ss58, netuid=fake_netuid, stake=100
+    )
+    fake_stake_info_3 = mocker.Mock(
+        hotkey_ss58=fake_hotkey_ss58, netuid=fake_netuid, stake=200
+    )
+
+    return_value = [fake_stake_info_1, fake_stake_info_2, fake_stake_info_3]
+
+    subtensor.get_stake_for_coldkey = mocker.patch.object(
+        subtensor, "get_stake_for_coldkey", return_value=return_value
+    )
+
+    # Call
+    result = subtensor.get_stake_for_coldkey_and_hotkey(
+        hotkey_ss58=fake_hotkey_ss58,
+        coldkey_ss58=fake_coldkey_ss58,
+        netuid=fake_netuid,
+        block=fake_block,
     )
-    if fake_value_result is not None:
-        spy_balance_from_rao.assert_called_once_with(fake_value_result)
-    else:
-        spy_balance_from_rao.assert_not_called()
-    assert result == fake_value_result
+
+    # Asserts
+    subtensor.get_stake_for_coldkey.assert_called_once_with(
+        fake_coldkey_ss58, fake_block
+    )
+    assert result == [fake_stake_info_2, fake_stake_info_3]
 
 
 def test_does_hotkey_exist_true(mocker, subtensor):
@@ -2714,6 +2747,7 @@ def test_add_stake_success(mocker, subtensor):
     fake_wallet = mocker.Mock()
     fake_hotkey_ss58 = "fake_hotkey"
     fake_amount = 10.0
+    fake_netuid = 123
 
     mock_add_stake_extrinsic = mocker.patch.object(
         subtensor_module, "add_stake_extrinsic"
@@ -2723,6 +2757,7 @@ def test_add_stake_success(mocker, subtensor):
     result = subtensor.add_stake(
         wallet=fake_wallet,
         hotkey_ss58=fake_hotkey_ss58,
+        netuid=fake_netuid,
         amount=fake_amount,
         wait_for_inclusion=True,
         wait_for_finalization=False,
@@ -2733,6 +2768,7 @@ def test_add_stake_success(mocker, subtensor):
         subtensor=subtensor,
         wallet=fake_wallet,
         hotkey_ss58=fake_hotkey_ss58,
+        netuid=fake_netuid,
         amount=fake_amount,
         wait_for_inclusion=True,
         wait_for_finalization=False,
@@ -2746,6 +2782,7 @@ def test_add_stake_multiple_success(mocker, subtensor):
     fake_wallet = mocker.Mock()
     fake_hotkey_ss58 = ["fake_hotkey"]
     fake_amount = [10.0]
+    fake_netuids = [1]
 
     mock_add_stake_multiple_extrinsic = mocker.patch.object(
         subtensor_module, "add_stake_multiple_extrinsic"
@@ -2755,6 +2792,7 @@ def test_add_stake_multiple_success(mocker, subtensor):
     result = subtensor.add_stake_multiple(
         wallet=fake_wallet,
         hotkey_ss58s=fake_hotkey_ss58,
+        netuids=fake_netuids,
         amounts=fake_amount,
         wait_for_inclusion=True,
         wait_for_finalization=False,
@@ -2765,6 +2803,7 @@ def test_add_stake_multiple_success(mocker, subtensor):
         subtensor=subtensor,
         wallet=fake_wallet,
         hotkey_ss58s=fake_hotkey_ss58,
+        netuids=fake_netuids,
         amounts=fake_amount,
         wait_for_inclusion=True,
         wait_for_finalization=False,
@@ -2778,13 +2817,15 @@ def test_unstake_success(mocker, subtensor):
     fake_wallet = mocker.Mock()
     fake_hotkey_ss58 = "hotkey_1"
     fake_amount = 10.0
+    fake_netuid = 123
 
     mock_unstake_extrinsic = mocker.patch.object(subtensor_module, "unstake_extrinsic")
 
     # Call
     result = subtensor.unstake(
         wallet=fake_wallet,
         hotkey_ss58=fake_hotkey_ss58,
+        netuid=fake_netuid,
         amount=fake_amount,
         wait_for_inclusion=True,
         wait_for_finalization=False,
@@ -2795,6 +2836,7 @@ def test_unstake_success(mocker, subtensor):
         subtensor=subtensor,
         wallet=fake_wallet,
         hotkey_ss58=fake_hotkey_ss58,
+        netuid=fake_netuid,
         amount=fake_amount,
         wait_for_inclusion=True,
         wait_for_finalization=False,
@@ -2808,6 +2850,7 @@ def test_unstake_multiple_success(mocker, subtensor):
     fake_wallet = mocker.Mock()
     fake_hotkeys = ["hotkey_1", "hotkey_2"]
     fake_amounts = [10.0, 20.0]
+    fake_netuids = [1, 2]
 
     mock_unstake_multiple_extrinsic = mocker.patch(
         "bittensor.core.subtensor.unstake_multiple_extrinsic", return_value=True
@@ -2817,6 +2860,7 @@ def test_unstake_multiple_success(mocker, subtensor):
     result = subtensor.unstake_multiple(
         wallet=fake_wallet,
         hotkey_ss58s=fake_hotkeys,
+        netuids=fake_netuids,
         amounts=fake_amounts,
         wait_for_inclusion=True,
         wait_for_finalization=False,
@@ -2827,6 +2871,7 @@ def test_unstake_multiple_success(mocker, subtensor):
         subtensor=subtensor,
         wallet=fake_wallet,
         hotkey_ss58s=fake_hotkeys,
+        netuids=fake_netuids,
         amounts=fake_amounts,
         wait_for_inclusion=True,
         wait_for_finalization=False,