[RAO] fix for unit test + refactoring

bittensor/core/async_subtensor.py

@@ -76,7 +76,7 @@
     u16_normalized_float,
     _decode_hex_identity_dict,
 )
-from bittensor.utils.balance import Balance, fixed_to_float, FixedPoint
+from bittensor.utils.balance import Balance, fixed_to_float
 from bittensor.utils.btlogging import logging
 from bittensor.utils.delegates_details import DelegatesDetails
 from bittensor.utils.weight_utils import generate_weight_hash
@@ -822,6 +822,8 @@ async def get_balance(
         )
         return Balance(balance["data"]["free"])
 
+    balance = get_balance
+
     async def get_balances(
         self,
         *addresses: str,
@@ -1501,7 +1503,7 @@ async def get_stake(
         block_hash = await self.determine_block_hash(block, block_hash, reuse_block)
 
         # Get alpha shares
-        alpha_shares: FixedPoint = await self.query_module(
+        alpha_shares = await self.query_module(
             module="SubtensorModule",
             name="Alpha",
             block_hash=block_hash,
@@ -1520,7 +1522,7 @@ async def get_stake(
         hotkey_alpha: int = getattr(hotkey_alpha_result, "value", 0)
 
         # Get total hotkey shares
-        hotkey_shares: FixedPoint = await self.query_module(
+        hotkey_shares = await self.query_module(
             module="SubtensorModule",
             name="TotalHotkeyShares",
             block_hash=block_hash,
@@ -1541,20 +1543,28 @@ async def get_stake(
     get_stake_for_coldkey_and_hotkey = get_stake
 
     async def get_stake_for_coldkey(
-        self, coldkey_ss58: str, block: Optional[int] = None
+        self,
+        coldkey_ss58: str,
+        block: Optional[int] = None,
+        block_hash: Optional[str] = None,
+        reuse_block: bool = False,
     ) -> 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.
+            block_hash (Optional[str]): The hash of the blockchain block number for the query.
+            reuse_block (bool): Whether to reuse the last-used block hash.
 
         Returns:
             Optional[list[StakeInfo]]: A list of StakeInfo objects, or ``None`` if no stake information is found.
         """
         encoded_coldkey = ss58_to_vec_u8(coldkey_ss58)
-        block_hash = await self.determine_block_hash(block)
+        block_hash = await self.determine_block_hash(
+            block=block, block_hash=block_hash, reuse_block=reuse_block
+        )
 
         hex_bytes_result = await self.query_runtime_api(
             runtime_api="StakeInfoRuntimeApi",
@@ -3455,7 +3465,7 @@ async def transfer(
         self,
         wallet: "Wallet",
         dest: str,
-        amount: Union["Balance", float],
+        amount: "Balance",
         transfer_all: bool = False,
         wait_for_inclusion: bool = True,
         wait_for_finalization: bool = False,
@@ -3506,6 +3516,7 @@ async def unstake(
             wallet (bittensor_wallet.wallet): The wallet associated with the neuron from which the stake is being
                 removed.
             hotkey_ss58 (Optional[str]): The ``SS58`` address of the hotkey account to unstake from.
+            netuid (Optional[int]): Subnet unique ID.
             amount (Balance): The amount of TAO to unstake. If not specified, unstakes all.
             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.
@@ -3543,6 +3554,7 @@ async def unstake_multiple(
             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]): Subnets unique IDs.
             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.

bittensor/core/extrinsics/asyncex/move_stake.py

@@ -5,11 +5,11 @@
 
 if TYPE_CHECKING:
     from bittensor_wallet import Wallet
-    from bittensor.core.subtensor import Subtensor
+    from bittensor.core.async_subtensor import AsyncSubtensor
 
 
 async def transfer_stake_extrinsic(
-    subtensor: "Subtensor",
+    subtensor: "AsyncSubtensor",
     wallet: "Wallet",
     destination_coldkey_ss58: str,
     hotkey_ss58: str,
@@ -19,6 +19,24 @@ async def transfer_stake_extrinsic(
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
 ) -> bool:
+    """
+    Transfers stake from one coldkey to another in the Bittensor network.
+
+    Args:
+        subtensor (AsyncSubtensor): The subtensor instance to interact with the blockchain.
+        wallet (Wallet): The wallet containing the coldkey to authorize the transfer.
+        destination_coldkey_ss58 (str): SS58 address of the destination coldkey.
+        hotkey_ss58 (str): SS58 address of the hotkey associated with the stake.
+        origin_netuid (int): Network UID of the origin subnet.
+        destination_netuid (int): Network UID of the destination subnet.
+        amount (Balance): The amount of stake to transfer as a `Balance` object.
+        wait_for_inclusion (bool): If True, waits for transaction inclusion in a block. Defaults to `True`.
+        wait_for_finalization (bool): If True, waits for transaction finalization. Defaults to `False`.
+
+    Returns:
+        bool: True if the transfer was successful, False otherwise.
+    """
+
     amount.set_unit(netuid=origin_netuid)
     # Verify ownership
     hotkey_owner = await subtensor.get_hotkey_owner(hotkey_ss58)
@@ -107,7 +125,7 @@ async def transfer_stake_extrinsic(
 
 
 async def swap_stake_extrinsic(
-    subtensor: "Subtensor",
+    subtensor: "AsyncSubtensor",
     wallet: "Wallet",
     hotkey_ss58: str,
     origin_netuid: int,
@@ -116,6 +134,22 @@ async def swap_stake_extrinsic(
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
 ) -> bool:
+    """
+    Swaps stake from one subnet to another for a given hotkey in the Bittensor network.
+
+    Args:
+        subtensor (AsyncSubtensor): The subtensor instance to interact with the blockchain.
+        wallet (Wallet): The wallet containing the coldkey to authorize the swap.
+        hotkey_ss58 (str): SS58 address of the hotkey associated with the stake.
+        origin_netuid (int): Network UID of the origin subnet.
+        destination_netuid (int): Network UID of the destination subnet.
+        amount (Balance): The amount of stake to swap as a `Balance` object.
+        wait_for_inclusion (bool): If True, waits for transaction inclusion in a block. Defaults to True.
+        wait_for_finalization (bool): If True, waits for transaction finalization. Defaults to False.
+
+    Returns:
+        bool: True if the swap was successful, False otherwise.
+    """
     amount.set_unit(netuid=origin_netuid)
     # Verify ownership
     hotkey_owner = await subtensor.get_hotkey_owner(hotkey_ss58)
@@ -203,7 +237,7 @@ async def swap_stake_extrinsic(
 
 
 async def move_stake_extrinsic(
-    subtensor: "Subtensor",
+    subtensor: "AsyncSubtensor",
     wallet: "Wallet",
     origin_hotkey: str,
     origin_netuid: int,
@@ -213,6 +247,23 @@ async def move_stake_extrinsic(
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
 ) -> bool:
+    """
+    Moves stake from one hotkey to another within subnets in the Bittensor network.
+
+    Args:
+        subtensor (Subtensor): The subtensor instance to interact with the blockchain.
+        wallet (Wallet): The wallet containing the coldkey to authorize the move.
+        origin_hotkey (str): SS58 address of the origin hotkey associated with the stake.
+        origin_netuid (int): Network UID of the origin subnet.
+        destination_hotkey (str): SS58 address of the destination hotkey.
+        destination_netuid (int): Network UID of the destination subnet.
+        amount (Balance): The amount of stake to move as a `Balance` object.
+        wait_for_inclusion (bool): If True, waits for transaction inclusion in a block. Defaults to True.
+        wait_for_finalization (bool): If True, waits for transaction finalization. Defaults to False.
+
+    Returns:
+        bool: True if the move was successful, False otherwise.
+    """
     amount.set_unit(netuid=origin_netuid)
     # Verify ownership of origin hotkey
     origin_owner = await subtensor.get_hotkey_owner(origin_hotkey)

bittensor/core/extrinsics/asyncex/staking.py

@@ -5,6 +5,7 @@
 from bittensor.utils import unlock_key
 from bittensor.utils.balance import Balance
 from bittensor.utils.btlogging import logging
+from bittensor.core.extrinsics.utils import get_old_stakes
 
 if TYPE_CHECKING:
     from bittensor_wallet import Wallet
@@ -184,27 +185,6 @@ async def add_stake_multiple_extrinsic(
         success: `True` if extrinsic was finalized or included in the block. `True` if any wallet was staked. If we did
             not wait for finalization/inclusion, the response is `True`.
     """
-
-    async def get_old_stakes() -> list[Balance]:
-        old_stakes = []
-        all_stakes = await subtensor.get_stake_for_coldkey(
-            coldkey_ss58=wallet.coldkeypub.ss58_address,
-        )
-        for hotkey_ss58, netuid in zip(hotkey_ss58s, netuids):
-            stake = next(
-                (
-                    stake.stake
-                    for stake in all_stakes
-                    if stake.hotkey_ss58 == hotkey_ss58
-                    and stake.coldkey_ss58 == wallet.coldkeypub.ss58_address
-                    and stake.netuid == netuid
-                ),
-                Balance.from_tao(0),  # Default to 0 balance if no match found
-            )
-            old_stakes.append(stake)
-
-        return old_stakes
-
     if not isinstance(hotkey_ss58s, list) or not all(
         isinstance(hotkey_ss58, str) for hotkey_ss58 in hotkey_ss58s
     ):
@@ -239,7 +219,13 @@ async def get_old_stakes() -> list[Balance]:
         f":satellite: [magenta]Syncing with chain:[/magenta] [blue]{subtensor.network}[/blue] [magenta]...[/magenta]"
     )
     block_hash = await subtensor.substrate.get_chain_head()
-    old_stakes: list[Balance] = await get_old_stakes()
+
+    all_stakes = await subtensor.get_stake_for_coldkey(
+        coldkey_ss58=wallet.coldkeypub.ss58_address, block_hash=block_hash
+    )
+    old_stakes: list[Balance] = get_old_stakes(
+        wallet=wallet, hotkey_ss58s=hotkey_ss58s, netuids=netuids, all_stakes=all_stakes
+    )
 
     # Remove existential balance to keep key alive.
     # Keys must maintain a balance of at least 1000 rao to stay alive.

bittensor/core/extrinsics/asyncex/unstaking.py

@@ -5,6 +5,7 @@
 from bittensor.utils import unlock_key
 from bittensor.utils.balance import Balance
 from bittensor.utils.btlogging import logging
+from bittensor.core.extrinsics.utils import get_old_stakes
 
 if TYPE_CHECKING:
     from bittensor_wallet import Wallet
@@ -164,27 +165,6 @@ async def unstake_multiple_extrinsic(
         success (bool): Flag is ``True`` if extrinsic was finalized or included in the block. Flag is ``True`` if any
             wallet was unstaked. If we did not wait for finalization / inclusion, the response is ``True``.
     """
-
-    async def get_old_stakes() -> list[Balance]:
-        old_stakes = []
-        all_stakes = await subtensor.get_stake_for_coldkey(
-            coldkey_ss58=wallet.coldkeypub.ss58_address,
-        )
-        for hotkey_ss58, netuid in zip(hotkey_ss58s, netuids):
-            stake = next(
-                (
-                    stake.stake
-                    for stake in all_stakes
-                    if stake.hotkey_ss58 == hotkey_ss58
-                    and stake.coldkey_ss58 == wallet.coldkeypub.ss58_address
-                    and stake.netuid == netuid
-                ),
-                Balance.from_tao(0),  # Default to 0 balance if no match found
-            )
-            old_stakes.append(stake)
-
-        return old_stakes
-
     if not isinstance(hotkey_ss58s, list) or not all(
         isinstance(hotkey_ss58, str) for hotkey_ss58 in hotkey_ss58s
     ):
@@ -223,10 +203,18 @@ async def get_old_stakes() -> list[Balance]:
     logging.info(
         f":satellite: [magenta]Syncing with chain:[/magenta] [blue]{subtensor.network}[/blue] [magenta]...[/magenta]"
     )
+
     block_hash = await subtensor.substrate.get_chain_head()
-    old_balance, old_stakes = await asyncio.gather(
+
+    all_stakes, old_balance = await asyncio.gather(
+        subtensor.get_stake_for_coldkey(
+            coldkey_ss58=wallet.coldkeypub.ss58_address, block_hash=block_hash
+        ),
         subtensor.get_balance(wallet.coldkeypub.ss58_address, block_hash=block_hash),
-        get_old_stakes(),
+    )
+
+    old_stakes: list[Balance] = get_old_stakes(
+        wallet=wallet, hotkey_ss58s=hotkey_ss58s, netuids=netuids, all_stakes=all_stakes
     )
 
     successful_unstakes = 0

bittensor/core/extrinsics/staking.py

@@ -2,6 +2,7 @@
 from typing import Optional, TYPE_CHECKING, Sequence
 
 from bittensor.core.errors import StakeError, NotRegisteredError
+from bittensor.core.extrinsics.utils import get_old_stakes
 from bittensor.utils import unlock_key
 from bittensor.utils.balance import Balance
 from bittensor.utils.btlogging import logging
@@ -27,6 +28,7 @@ def add_stake_extrinsic(
         subtensor: the Subtensor object to use
         wallet: Bittensor wallet object.
         hotkey_ss58: The `ss58` address of the hotkey account to stake to defaults to the wallet's hotkey.
+        netuid (Optional[int]): Subnet unique ID.
         amount: Amount to stake as Bittensor balance, `None` if staking all.
         wait_for_inclusion: 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.
@@ -174,26 +176,6 @@ def add_stake_multiple_extrinsic(
             not wait for finalization/inclusion, the response is `True`.
     """
 
-    def get_old_stakes() -> list[Balance]:
-        old_stakes = []
-        all_stakes = subtensor.get_stake_for_coldkey(
-            coldkey_ss58=wallet.coldkeypub.ss58_address,
-        )
-        for hotkey_ss58, netuid in zip(hotkey_ss58s, netuids):
-            stake = next(
-                (
-                    stake.stake
-                    for stake in all_stakes
-                    if stake.hotkey_ss58 == hotkey_ss58
-                    and stake.coldkey_ss58 == wallet.coldkeypub.ss58_address
-                    and stake.netuid == netuid
-                ),
-                Balance.from_tao(0),  # Default to 0 balance if no match found
-            )
-            old_stakes.append(stake)
-
-        return old_stakes
-
     if not isinstance(hotkey_ss58s, list) or not all(
         isinstance(hotkey_ss58, str) for hotkey_ss58 in hotkey_ss58s
     ):
@@ -209,6 +191,7 @@ def get_old_stakes() -> list[Balance]:
         raise ValueError("netuids must be a list of the same length as hotkey_ss58s")
 
     new_amounts: Sequence[Optional[Balance]]
+
     if amounts is None:
         new_amounts = [None] * len(hotkey_ss58s)
     else:
@@ -228,7 +211,12 @@ def get_old_stakes() -> list[Balance]:
         f":satellite: [magenta]Syncing with chain:[/magenta] [blue]{subtensor.network}[/blue] [magenta]...[/magenta]"
     )
     block = subtensor.get_current_block()
-    old_stakes: list[Balance] = get_old_stakes()
+    all_stakes = subtensor.get_stake_for_coldkey(
+        coldkey_ss58=wallet.coldkeypub.ss58_address,
+    )
+    old_stakes: list[Balance] = get_old_stakes(
+        wallet=wallet, hotkey_ss58s=hotkey_ss58s, netuids=netuids, all_stakes=all_stakes
+    )
 
     # Remove existential balance to keep key alive.
     # Keys must maintain a balance of at least 1000 rao to stay alive.