diff --git a/bittensor/core/async_subtensor.py b/bittensor/core/async_subtensor.py index 02cc910f2a..66a28e2c33 100644 --- a/bittensor/core/async_subtensor.py +++ b/bittensor/core/async_subtensor.py @@ -2810,6 +2810,9 @@ async def add_stake( amount: Optional[Balance] = None, wait_for_inclusion: bool = True, wait_for_finalization: bool = False, + safe_staking: bool = False, + allow_partial_stake: bool = False, + rate_threshold: float = 0.005, ) -> bool: """ Adds the specified amount of stake to a neuron identified by the hotkey ``SS58`` address. @@ -2823,12 +2826,20 @@ async def add_stake( amount (Balance): The amount of TAO to stake. 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. + safe_staking (bool): If true, enables price safety checks to protect against fluctuating prices. The stake + will only execute if the price change doesn't exceed the rate threshold. Default is False. + allow_partial_stake (bool): If true and safe_staking is enabled, allows partial staking when + the full amount would exceed the price threshold. If false, the entire stake fails if it would + exceed the threshold. Default is False. + rate_threshold (float): The maximum allowed price change ratio when staking. For example, + 0.005 = 0.5% maximum price increase. Only used when safe_staking is True. Default is 0.005. Returns: bool: ``True`` if the staking is successful, False otherwise. - This function enables neurons to increase their stake in the network, enhancing their influence and potential - rewards in line with Bittensor's consensus and reward mechanisms. + This function enables neurons to increase their stake in the network, enhancing their influence and potential. + When safe_staking is enabled, it provides protection against price fluctuations during the time stake is + executed and the time it is actually processed by the chain. """ amount = check_and_convert_to_balance(amount) return await add_stake_extrinsic( @@ -2839,6 +2850,9 @@ async def add_stake( amount=amount, wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, + safe_staking=safe_staking, + allow_partial_stake=allow_partial_stake, + rate_threshold=rate_threshold, ) async def add_stake_multiple( @@ -3460,6 +3474,9 @@ async def swap_stake( amount: Balance, wait_for_inclusion: bool = True, wait_for_finalization: bool = False, + safe_staking: bool = False, + allow_partial_stake: bool = False, + rate_threshold: float = 0.005, ) -> bool: """ Moves stake between subnets while keeping the same coldkey-hotkey pair ownership. @@ -3473,9 +3490,25 @@ async def swap_stake( amount (Union[Balance, float]): The amount to swap. 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. + safe_staking (bool): If true, enables price safety checks to protect against fluctuating prices. The swap + will only execute if the price ratio between subnets doesn't exceed the rate threshold. + Default is False. + allow_partial_stake (bool): If true and safe_staking is enabled, allows partial stake swaps when + the full amount would exceed the price threshold. If false, the entire swap fails if it would + exceed the threshold. Default is False. + rate_threshold (float): The maximum allowed increase in the price ratio between subnets + (origin_price/destination_price). For example, 0.005 = 0.5% maximum increase. Only used + when safe_staking is True. Default is 0.005. Returns: success (bool): True if the extrinsic was successful. + + The price ratio for swap_stake in safe mode is calculated as: origin_subnet_price / destination_subnet_price + When safe_staking is enabled, the swap will only execute if: + - With allow_partial_stake=False: The entire swap amount can be executed without the price ratio + increasing more than rate_threshold + - With allow_partial_stake=True: A partial amount will be swapped up to the point where the + price ratio would increase by rate_threshold """ amount = check_and_convert_to_balance(amount) return await swap_stake_extrinsic( @@ -3487,6 +3520,9 @@ async def swap_stake( amount=amount, wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, + safe_staking=safe_staking, + allow_partial_stake=allow_partial_stake, + rate_threshold=rate_threshold, ) async def transfer_stake( @@ -3575,6 +3611,9 @@ async def unstake( amount: Optional[Balance] = None, wait_for_inclusion: bool = True, wait_for_finalization: bool = False, + safe_staking: bool = False, + allow_partial_stake: bool = False, + rate_threshold: float = 0.005, ) -> bool: """ Removes a specified amount of stake from a single hotkey account. This function is critical for adjusting @@ -3584,10 +3623,17 @@ 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. + netuid (Optional[int]): The unique identifier of the subnet. 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. + safe_staking (bool): If true, enables price safety checks to protect against fluctuating prices. The unstake + will only execute if the price change doesn't exceed the rate threshold. Default is False. + allow_partial_stake (bool): If true and safe_staking is enabled, allows partial unstaking when + the full amount would exceed the price threshold. If false, the entire unstake fails if it would + exceed the threshold. Default is False. + rate_threshold (float): The maximum allowed price change ratio when unstaking. For example, + 0.005 = 0.5% maximum price decrease. Only used when safe_staking is True. Default is 0.005. Returns: bool: ``True`` if the unstaking process is successful, False otherwise. @@ -3604,6 +3650,9 @@ async def unstake( amount=amount, wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, + safe_staking=safe_staking, + allow_partial_stake=allow_partial_stake, + rate_threshold=rate_threshold, ) async def unstake_multiple( diff --git a/bittensor/core/chain_data/dynamic_info.py b/bittensor/core/chain_data/dynamic_info.py index 9f2739560b..cde40a0b79 100644 --- a/bittensor/core/chain_data/dynamic_info.py +++ b/bittensor/core/chain_data/dynamic_info.py @@ -77,9 +77,9 @@ def _from_dict(cls, decoded: dict) -> "DynamicInfo": price = ( Balance.from_tao(1.0) if netuid == 0 - else Balance.from_tao(tao_in.tao / alpha_in.tao) + else Balance.from_tao(tao_in.tao / alpha_in.tao).set_unit(netuid) if alpha_in.tao > 0 - else Balance.from_tao(1) + else Balance.from_tao(1).set_unit(netuid) ) # Root always has 1-1 price if decoded.get("subnet_identity"): diff --git a/bittensor/core/extrinsics/asyncex/move_stake.py b/bittensor/core/extrinsics/asyncex/move_stake.py index 4477f5d960..4c1a3ee0f7 100644 --- a/bittensor/core/extrinsics/asyncex/move_stake.py +++ b/bittensor/core/extrinsics/asyncex/move_stake.py @@ -160,6 +160,9 @@ async def swap_stake_extrinsic( amount: Balance, wait_for_inclusion: bool = True, wait_for_finalization: bool = False, + safe_staking: bool = False, + allow_partial_stake: bool = False, + rate_threshold: float = 0.005, ) -> bool: """ Swaps stake from one subnet to another for a given hotkey in the Bittensor network. @@ -173,6 +176,9 @@ async def swap_stake_extrinsic( 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. + safe_staking (bool): If true, enables price safety checks to protect against price impact. + allow_partial_stake (bool): If true, allows partial stake swaps when the full amount would exceed the price threshold. + rate_threshold (float): Maximum allowed increase in price ratio (0.005 = 0.5%). Returns: bool: True if the swap was successful, False otherwise. @@ -205,20 +211,47 @@ async def swap_stake_extrinsic( return False try: - logging.info( - f"Swapping stake for hotkey [blue]{hotkey_ss58}[/blue]\n" - f"Amount: [green]{amount}[/green] from netuid [yellow]{origin_netuid}[/yellow] to netuid " - f"[yellow]{destination_netuid}[/yellow]" - ) + call_params = { + "hotkey": hotkey_ss58, + "origin_netuid": origin_netuid, + "destination_netuid": destination_netuid, + "alpha_amount": amount.rao, + } + + if safe_staking: + origin_pool, destination_pool = await asyncio.gather( + subtensor.subnet(netuid=origin_netuid), + subtensor.subnet(netuid=destination_netuid), + ) + swap_rate_ratio = origin_pool.price.rao / destination_pool.price.rao + swap_rate_ratio_with_tolerance = swap_rate_ratio * (1 + rate_threshold) + + logging.info( + f"Swapping stake with safety for hotkey [blue]{hotkey_ss58}[/blue]\n" + f"Amount: [green]{amount}[/green] from netuid [green]{origin_netuid}[/green] to netuid " + f"[green]{destination_netuid}[/green]\n" + f"Current price ratio: [green]{swap_rate_ratio:.4f}[/green], " + f"Ratio with tolerance: [green]{swap_rate_ratio_with_tolerance:.4f}[/green]" + ) + call_params.update( + { + "limit_price": swap_rate_ratio_with_tolerance, + "allow_partial": allow_partial_stake, + } + ) + call_function = "swap_stake_limit" + else: + logging.info( + f"Swapping stake for hotkey [blue]{hotkey_ss58}[/blue]\n" + f"Amount: [green]{amount}[/green] from netuid [green]{origin_netuid}[/green] to netuid " + f"[green]{destination_netuid}[/green]" + ) + call_function = "swap_stake" + call = await subtensor.substrate.compose_call( call_module="SubtensorModule", - call_function="swap_stake", - call_params={ - "hotkey": hotkey_ss58, - "origin_netuid": origin_netuid, - "destination_netuid": destination_netuid, - "alpha_amount": amount.rao, - }, + call_function=call_function, + call_params=call_params, ) success, err_msg = await subtensor.sign_and_send_extrinsic( @@ -253,7 +286,12 @@ async def swap_stake_extrinsic( return True else: - logging.error(f":cross_mark: [red]Failed[/red]: {err_msg}") + if safe_staking and "Custom error: 8" in err_msg: + logging.error( + ":cross_mark: [red]Failed[/red]: Price ratio exceeded tolerance limit. Either increase price tolerance or enable partial staking." + ) + else: + logging.error(f":cross_mark: [red]Failed[/red]: {err_msg}") return False except Exception as e: diff --git a/bittensor/core/extrinsics/asyncex/staking.py b/bittensor/core/extrinsics/asyncex/staking.py index 36e7d3f6d5..8212ff285e 100644 --- a/bittensor/core/extrinsics/asyncex/staking.py +++ b/bittensor/core/extrinsics/asyncex/staking.py @@ -21,6 +21,9 @@ async def add_stake_extrinsic( amount: Optional[Balance] = None, wait_for_inclusion: bool = True, wait_for_finalization: bool = False, + safe_staking: bool = False, + allow_partial_stake: bool = False, + rate_threshold: float = 0.005, ) -> bool: """ Adds the specified amount of stake to passed hotkey `uid`. @@ -36,6 +39,9 @@ async def add_stake_extrinsic( `False` if the extrinsic fails to enter the block within the timeout. wait_for_finalization: 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. + safe_staking: If set, uses safe staking logic + allow_partial_stake: If set, allows partial stake + rate_threshold: The rate threshold for safe staking Returns: success: Flag is `True` if extrinsic was finalized or included in the block. If we did not wait for @@ -97,19 +103,48 @@ async def add_stake_extrinsic( return False try: - logging.info( - f":satellite: [magenta]Staking to:[/magenta] " - f"[blue]netuid: {netuid}, amount: {staking_balance} " - f"on {subtensor.network}[/blue] [magenta]...[/magenta]" - ) + call_params = { + "hotkey": hotkey_ss58, + "netuid": netuid, + "amount_staked": staking_balance.rao, + } + + if safe_staking: + pool = await subtensor.subnet(netuid=netuid) + base_price = pool.price.rao + price_with_tolerance = base_price * (1 + rate_threshold) + call_params.update( + { + "limit_price": price_with_tolerance, + "allow_partial": allow_partial_stake, + } + ) + call_function = "add_stake_limit" + + # For logging + base_rate = pool.price.tao + rate_with_tolerance = base_rate * (1 + rate_threshold) + logging.info( + f":satellite: [magenta]Safe Staking to:[/magenta] " + f"[blue]netuid: [green]{netuid}[/green], amount: [green]{staking_balance}[/green], " + f"tolerance percentage: [green]{rate_threshold*100}%[/green], " + f"price limit: [green]{rate_with_tolerance}[/green], " + f"original price: [green]{base_rate}[/green], " + f"with partial stake: [green]{allow_partial_stake}[/green] " + f"on [blue]{subtensor.network}[/blue][/magenta]...[/magenta]" + ) + else: + logging.info( + f":satellite: [magenta]Staking to:[/magenta] " + f"[blue]netuid: [green]{netuid}[/green], amount: [green]{staking_balance}[/green] " + f"on [blue]{subtensor.network}[/blue][magenta]...[/magenta]" + ) + call_function = "add_stake" + call = await subtensor.substrate.compose_call( call_module="SubtensorModule", - call_function="add_stake", - call_params={ - "hotkey": hotkey_ss58, - "amount_staked": staking_balance.rao, - "netuid": netuid, - }, + call_function=call_function, + call_params=call_params, ) staking_response, err_msg = await subtensor.sign_and_send_extrinsic( call, @@ -152,7 +187,12 @@ async def add_stake_extrinsic( ) return True else: - logging.error(f":cross_mark: [red]Failed: {err_msg}.[/red]") + if safe_staking and "Custom error: 8" in err_msg: + logging.error( + ":cross_mark: [red]Failed[/red]: Price exceeded tolerance limit. Either increase price tolerance or enable partial staking." + ) + else: + logging.error(f":cross_mark: [red]Failed: {err_msg}.[/red]") return False except NotRegisteredError: diff --git a/bittensor/core/extrinsics/asyncex/unstaking.py b/bittensor/core/extrinsics/asyncex/unstaking.py index 7ba1d120a2..f8844f59e8 100644 --- a/bittensor/core/extrinsics/asyncex/unstaking.py +++ b/bittensor/core/extrinsics/asyncex/unstaking.py @@ -20,6 +20,9 @@ async def unstake_extrinsic( amount: Optional[Balance] = None, wait_for_inclusion: bool = True, wait_for_finalization: bool = False, + safe_staking: bool = False, + allow_partial_stake: bool = False, + rate_threshold: float = 0.005, ) -> bool: """Removes stake into the wallet coldkey from the specified hotkey ``uid``. @@ -34,6 +37,9 @@ async def unstake_extrinsic( 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. + safe_staking: If true, enables price safety checks + allow_partial_stake: If true, allows partial unstaking if price threshold exceeded + rate_threshold: Maximum allowed price decrease percentage (0.005 = 0.5%) Returns: success (bool): Flag is ``True`` if extrinsic was finalized or included in the block. If we did not wait for @@ -83,19 +89,49 @@ async def unstake_extrinsic( return False try: - logging.info( - f"Unstaking [blue]{unstaking_balance}[/blue] from hotkey: [magenta]{hotkey_ss58}[/magenta] on netuid: " - f"[blue]{netuid}[/blue]" - ) + call_params = { + "hotkey": hotkey_ss58, + "netuid": netuid, + "amount_unstaked": unstaking_balance.rao, + } + if safe_staking: + pool = await subtensor.subnet(netuid=netuid) + base_price = pool.price.rao + price_with_tolerance = base_price * (1 - rate_threshold) + + # For logging + base_rate = pool.price.tao + rate_with_tolerance = base_rate * (1 - rate_threshold) + + logging.info( + f":satellite: [magenta]Safe Unstaking from:[/magenta] " + f"netuid: [green]{netuid}[/green], amount: [green]{unstaking_balance}[/green], " + f"tolerance percentage: [green]{rate_threshold*100}%[/green], " + f"price limit: [green]{rate_with_tolerance}[/green], " + f"original price: [green]{base_rate}[/green], " + f"with partial unstake: [green]{allow_partial_stake}[/green] " + f"on [blue]{subtensor.network}[/blue][magenta]...[/magenta]" + ) + + call_params.update( + { + "limit_price": price_with_tolerance, + "allow_partial": allow_partial_stake, + } + ) + call_function = "remove_stake_limit" + else: + logging.info( + f":satellite: [magenta]Unstaking from:[/magenta] " + f"netuid: [green]{netuid}[/green], amount: [green]{unstaking_balance}[/green] " + f"on [blue]{subtensor.network}[/blue][magenta]...[/magenta]" + ) + call_function = "remove_stake" call = await subtensor.substrate.compose_call( call_module="SubtensorModule", - call_function="remove_stake", - call_params={ - "hotkey": hotkey_ss58, - "amount_unstaked": unstaking_balance.rao, - "netuid": netuid, - }, + call_function=call_function, + call_params=call_params, ) staking_response, err_msg = await subtensor.sign_and_send_extrinsic( call, @@ -138,7 +174,12 @@ async def unstake_extrinsic( ) return True else: - logging.error(f":cross_mark: [red]Failed: {err_msg}.[/red]") + if safe_staking and "Custom error: 8" in err_msg: + logging.error( + ":cross_mark: [red]Failed[/red]: Price exceeded tolerance limit. Either increase price tolerance or enable partial staking." + ) + else: + logging.error(f":cross_mark: [red]Failed: {err_msg}.[/red]") return False except NotRegisteredError: diff --git a/bittensor/core/extrinsics/move_stake.py b/bittensor/core/extrinsics/move_stake.py index 77cd4a986a..8a1d3a7281 100644 --- a/bittensor/core/extrinsics/move_stake.py +++ b/bittensor/core/extrinsics/move_stake.py @@ -157,6 +157,9 @@ def swap_stake_extrinsic( amount: Optional[Balance] = None, wait_for_inclusion: bool = True, wait_for_finalization: bool = False, + safe_staking: bool = False, + allow_partial_stake: bool = False, + rate_threshold: float = 0.005, ) -> bool: """ Moves stake between subnets while keeping the same coldkey-hotkey pair ownership. @@ -170,6 +173,9 @@ def swap_stake_extrinsic( amount (Union[Balance, float]): Amount to swap. wait_for_inclusion (bool): If true, waits for inclusion before returning. wait_for_finalization (bool): If true, waits for finalization before returning. + safe_staking (bool): If true, enables price safety checks to protect against price impact. + allow_partial_stake (bool): If true, allows partial stake swaps when the full amount would exceed the price threshold. + rate_threshold (float): Maximum allowed increase in price ratio (0.005 = 0.5%). Returns: success (bool): True if the swap was successful. @@ -203,20 +209,45 @@ def swap_stake_extrinsic( return False try: - logging.info( - f"Swapping stake for hotkey [blue]{hotkey_ss58}[/blue]\n" - f"Amount: [green]{amount}[/green] from netuid [yellow]{origin_netuid}[/yellow] to netuid " - f"[yellow]{destination_netuid}[/yellow]" - ) + call_params = { + "hotkey": hotkey_ss58, + "origin_netuid": origin_netuid, + "destination_netuid": destination_netuid, + "alpha_amount": amount.rao, + } + + if safe_staking: + origin_pool = subtensor.subnet(netuid=origin_netuid) + destination_pool = subtensor.subnet(netuid=destination_netuid) + swap_rate_ratio = origin_pool.price.rao / destination_pool.price.rao + swap_rate_ratio_with_tolerance = swap_rate_ratio * (1 + rate_threshold) + + logging.info( + f"Swapping stake with safety for hotkey [blue]{hotkey_ss58}[/blue]\n" + f"Amount: [green]{amount}[/green] from netuid [green]{origin_netuid}[/green] to netuid " + f"[green]{destination_netuid}[/green]\n" + f"Current price ratio: [green]{swap_rate_ratio:.4f}[/green], " + f"Ratio with tolerance: [green]{swap_rate_ratio_with_tolerance:.4f}[/green]" + ) + call_params.update( + { + "limit_price": swap_rate_ratio_with_tolerance, + "allow_partial": allow_partial_stake, + } + ) + call_function = "swap_stake_limit" + else: + logging.info( + f"Swapping stake for hotkey [blue]{hotkey_ss58}[/blue]\n" + f"Amount: [green]{amount}[/green] from netuid [green]{origin_netuid}[/green] to netuid " + f"[green]{destination_netuid}[/green]" + ) + call_function = "swap_stake" + call = subtensor.substrate.compose_call( call_module="SubtensorModule", - call_function="swap_stake", - call_params={ - "hotkey": hotkey_ss58, - "origin_netuid": origin_netuid, - "destination_netuid": destination_netuid, - "alpha_amount": amount.rao, - }, + call_function=call_function, + call_params=call_params, ) success, err_msg = subtensor.sign_and_send_extrinsic( @@ -251,7 +282,12 @@ def swap_stake_extrinsic( return True else: - logging.error(f":cross_mark: [red]Failed[/red]: {err_msg}") + if safe_staking and "Custom error: 8" in err_msg: + logging.error( + ":cross_mark: [red]Failed[/red]: Price ratio exceeded tolerance limit. Either increase price tolerance or enable partial staking." + ) + else: + logging.error(f":cross_mark: [red]Failed[/red]: {err_msg}") return False except Exception as e: diff --git a/bittensor/core/extrinsics/staking.py b/bittensor/core/extrinsics/staking.py index 9872eee0f4..6207a32686 100644 --- a/bittensor/core/extrinsics/staking.py +++ b/bittensor/core/extrinsics/staking.py @@ -19,6 +19,9 @@ def add_stake_extrinsic( amount: Optional[Balance] = None, wait_for_inclusion: bool = True, wait_for_finalization: bool = False, + safe_staking: bool = False, + allow_partial_stake: bool = False, + rate_threshold: float = 0.005, ) -> bool: """ Adds the specified amount of stake to passed hotkey `uid`. @@ -33,6 +36,9 @@ def add_stake_extrinsic( `False` if the extrinsic fails to enter the block within the timeout. wait_for_finalization: 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. + safe_staking (bool): If true, enables price safety checks + allow_partial_stake (bool): If true, allows partial unstaking if price threshold exceeded + rate_threshold (float): Maximum allowed price increase percentage (0.005 = 0.5%) Returns: success: Flag is `True` if extrinsic was finalized or included in the block. If we did not wait for @@ -91,20 +97,52 @@ def add_stake_extrinsic( return False try: - logging.info( - f":satellite: [magenta]Staking to:[/magenta] " - f"[blue]netuid: {netuid}, amount: {staking_balance} " - f"on {subtensor.network}[/blue] [magenta]...[/magenta]" - ) + call_params = { + "hotkey": hotkey_ss58, + "netuid": netuid, + "amount_staked": staking_balance.rao, + } + + if safe_staking: + pool = subtensor.subnet(netuid=netuid) + base_price = pool.price.rao + price_with_tolerance = base_price * (1 + rate_threshold) + + # For logging + base_rate = pool.price.tao + rate_with_tolerance = base_rate * (1 + rate_threshold) + + logging.info( + f":satellite: [magenta]Safe Staking to:[/magenta] " + f"[blue]netuid: [green]{netuid}[/green], amount: [green]{staking_balance}[/green], " + f"tolerance percentage: [green]{rate_threshold*100}%[/green], " + f"price limit: [green]{rate_with_tolerance}[/green], " + f"original price: [green]{base_rate}[/green], " + f"with partial stake: [green]{allow_partial_stake}[/green] " + f"on [blue]{subtensor.network}[/blue][/magenta]...[/magenta]" + ) + + call_params.update( + { + "limit_price": price_with_tolerance, + "allow_partial": allow_partial_stake, + } + ) + call_function = "add_stake_limit" + else: + logging.info( + f":satellite: [magenta]Staking to:[/magenta] " + f"[blue]netuid: [green]{netuid}[/green], amount: [green]{staking_balance}[/green] " + f"on [blue]{subtensor.network}[/blue][magenta]...[/magenta]" + ) + call_function = "add_stake" + call = subtensor.substrate.compose_call( call_module="SubtensorModule", - call_function="add_stake", - call_params={ - "hotkey": hotkey_ss58, - "amount_staked": staking_balance.rao, - "netuid": netuid, - }, + call_function=call_function, + call_params=call_params, ) + staking_response, err_msg = subtensor.sign_and_send_extrinsic( call, wallet, @@ -143,7 +181,12 @@ def add_stake_extrinsic( ) return True else: - logging.error(":cross_mark: [red]Failed[/red]: Error unknown.") + if safe_staking and "Custom error: 8" in err_msg: + logging.error( + ":cross_mark: [red]Failed[/red]: Price exceeded tolerance limit. Either increase price tolerance or enable partial staking." + ) + else: + logging.error(f":cross_mark: [red]Failed: {err_msg}.[/red]") return False # TODO I don't think these are used. Maybe should just catch SubstrateRequestException? diff --git a/bittensor/core/extrinsics/unstaking.py b/bittensor/core/extrinsics/unstaking.py index eb178f56f9..d0811fbebb 100644 --- a/bittensor/core/extrinsics/unstaking.py +++ b/bittensor/core/extrinsics/unstaking.py @@ -19,6 +19,9 @@ def unstake_extrinsic( amount: Optional[Balance] = None, wait_for_inclusion: bool = True, wait_for_finalization: bool = False, + safe_staking: bool = False, + allow_partial_stake: bool = False, + rate_threshold: float = 0.005, ) -> bool: """Removes stake into the wallet coldkey from the specified hotkey ``uid``. @@ -33,6 +36,9 @@ def unstake_extrinsic( 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. + safe_staking: If true, enables price safety checks + allow_partial_stake: If true, allows partial unstaking if price threshold exceeded + rate_threshold: Maximum allowed price decrease percentage (0.005 = 0.5%) Returns: success (bool): Flag is ``True`` if extrinsic was finalized or included in the block. If we did not wait for @@ -79,18 +85,52 @@ def unstake_extrinsic( return False try: - logging.info( - f"Unstaking [blue]{unstaking_balance}[/blue] from [magenta]{hotkey_ss58}[/magenta] on [blue]{netuid}[/blue]" - ) + call_params = { + "hotkey": hotkey_ss58, + "netuid": netuid, + "amount_unstaked": unstaking_balance.rao, + } + + if safe_staking: + pool = subtensor.subnet(netuid=netuid) + base_price = pool.price.rao + price_with_tolerance = base_price * (1 - rate_threshold) + + # For logging + base_rate = pool.price.tao + rate_with_tolerance = base_rate * (1 - rate_threshold) + + logging.info( + f":satellite: [magenta]Safe Unstaking from:[/magenta] " + f"netuid: [green]{netuid}[/green], amount: [green]{unstaking_balance}[/green], " + f"tolerance percentage: [green]{rate_threshold*100}%[/green], " + f"price limit: [green]{rate_with_tolerance}[/green], " + f"original price: [green]{base_rate}[/green], " + f"with partial unstake: [green]{allow_partial_stake}[/green] " + f"on [blue]{subtensor.network}[/blue][magenta]...[/magenta]" + ) + + call_params.update( + { + "limit_price": price_with_tolerance, + "allow_partial": allow_partial_stake, + } + ) + call_function = "remove_stake_limit" + else: + logging.info( + f":satellite: [magenta]Unstaking from:[/magenta] " + f"netuid: [green]{netuid}[/green], amount: [green]{unstaking_balance}[/green] " + f"on [blue]{subtensor.network}[/blue][magenta]...[/magenta]" + ) + call_function = "remove_stake" + call = subtensor.substrate.compose_call( call_module="SubtensorModule", - call_function="remove_stake", - call_params={ - "hotkey": hotkey_ss58, - "amount_unstaked": unstaking_balance.rao, - "netuid": netuid, - }, + call_function=call_function, + call_params=call_params, ) + staking_response, err_msg = subtensor.sign_and_send_extrinsic( call, wallet, @@ -130,7 +170,12 @@ def unstake_extrinsic( ) return True else: - logging.error(f":cross_mark: [red]Failed: {err_msg}.[/red]") + if safe_staking and "Custom error: 8" in err_msg: + logging.error( + ":cross_mark: [red]Failed[/red]: Price exceeded tolerance limit. Either increase price tolerance or enable partial staking." + ) + else: + logging.error(f":cross_mark: [red]Failed: {err_msg}.[/red]") return False except NotRegisteredError: diff --git a/bittensor/core/subtensor.py b/bittensor/core/subtensor.py index 406074bb12..8496c911b5 100644 --- a/bittensor/core/subtensor.py +++ b/bittensor/core/subtensor.py @@ -2129,6 +2129,9 @@ def add_stake( amount: Optional[Balance] = None, wait_for_inclusion: bool = True, wait_for_finalization: bool = False, + safe_staking: bool = False, + allow_partial_stake: bool = False, + rate_threshold: float = 0.005, ) -> bool: """ Adds the specified amount of stake to a neuron identified by the hotkey ``SS58`` address. @@ -2142,12 +2145,21 @@ def add_stake( amount (Balance): The amount of TAO to stake. 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. + safe_staking (bool): If true, enables price safety checks to protect against fluctuating prices. The stake + will only execute if the price change doesn't exceed the rate threshold. Default is False. + allow_partial_stake (bool): If true and safe_staking is enabled, allows partial staking when + the full amount would exceed the price threshold. If false, the entire stake fails if it would + exceed the threshold. Default is False. + rate_threshold (float): The maximum allowed price change ratio when staking. For example, + 0.005 = 0.5% maximum price increase. Only used when safe_staking is True. Default is 0.005. Returns: - bool: ``True`` if the staking is successful, False otherwise. + bool: True if the staking is successful, False otherwise. This function enables neurons to increase their stake in the network, enhancing their influence and potential rewards in line with Bittensor's consensus and reward mechanisms. + When safe_staking is enabled, it provides protection against price fluctuations during the time stake is + executed and the time it is actually processed by the chain. """ amount = check_and_convert_to_balance(amount) return add_stake_extrinsic( @@ -2158,6 +2170,9 @@ def add_stake( amount=amount, wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, + safe_staking=safe_staking, + allow_partial_stake=allow_partial_stake, + rate_threshold=rate_threshold, ) def add_stake_multiple( @@ -2760,6 +2775,9 @@ def swap_stake( amount: Balance, wait_for_inclusion: bool = True, wait_for_finalization: bool = False, + safe_staking: bool = False, + allow_partial_stake: bool = False, + rate_threshold: float = 0.005, ) -> bool: """ Moves stake between subnets while keeping the same coldkey-hotkey pair ownership. @@ -2773,9 +2791,26 @@ def swap_stake( amount (Union[Balance, float]): The amount to swap. 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. + safe_staking (bool): If true, enables price safety checks to protect against fluctuating prices. The swap + will only execute if the price ratio between subnets doesn't exceed the rate threshold. + Default is False. + allow_partial_stake (bool): If true and safe_staking is enabled, allows partial stake swaps when + the full amount would exceed the price threshold. If false, the entire swap fails if it would + exceed the threshold. Default is False. + rate_threshold (float): The maximum allowed increase in the price ratio between subnets + (origin_price/destination_price). For example, 0.005 = 0.5% maximum increase. Only used + when safe_staking is True. Default is 0.005. + Returns: success (bool): True if the extrinsic was successful. + + The price ratio for swap_stake in safe mode is calculated as: origin_subnet_price / destination_subnet_price + When safe_staking is enabled, the swap will only execute if: + - With allow_partial_stake=False: The entire swap amount can be executed without the price ratio + increasing more than rate_threshold + - With allow_partial_stake=True: A partial amount will be swapped up to the point where the + price ratio would increase by rate_threshold """ amount = check_and_convert_to_balance(amount) return swap_stake_extrinsic( @@ -2787,6 +2822,9 @@ def swap_stake( amount=amount, wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, + safe_staking=safe_staking, + allow_partial_stake=allow_partial_stake, + rate_threshold=rate_threshold, ) def transfer( @@ -2875,6 +2913,9 @@ def unstake( amount: Optional[Balance] = None, wait_for_inclusion: bool = True, wait_for_finalization: bool = False, + safe_staking: bool = False, + allow_partial_stake: bool = False, + rate_threshold: float = 0.005, ) -> bool: """ Removes a specified amount of stake from a single hotkey account. This function is critical for adjusting @@ -2888,12 +2929,20 @@ def unstake( 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. + safe_staking (bool): If true, enables price safety checks to protect against fluctuating prices. The unstake + will only execute if the price change doesn't exceed the rate threshold. Default is False. + allow_partial_stake (bool): If true and safe_staking is enabled, allows partial unstaking when + the full amount would exceed the price threshold. If false, the entire unstake fails if it would + exceed the threshold. Default is False. + rate_threshold (float): The maximum allowed price change ratio when unstaking. For example, + 0.005 = 0.5% maximum price decrease. Only used when safe_staking is True. Default is 0.005. Returns: bool: ``True`` if the unstaking process is successful, False otherwise. This function supports flexible stake management, allowing neurons to adjust their network participation and - potential reward accruals. + potential reward accruals. When safe_staking is enabled, it provides protection against price fluctuations + during the time unstake is executed and the time it is actually processed by the chain. """ amount = check_and_convert_to_balance(amount) return unstake_extrinsic( @@ -2904,6 +2953,9 @@ def unstake( amount=amount, wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, + safe_staking=safe_staking, + allow_partial_stake=allow_partial_stake, + rate_threshold=rate_threshold, ) def unstake_multiple( diff --git a/tests/e2e_tests/test_staking.py b/tests/e2e_tests/test_staking.py index 64d61f466d..105c52595e 100644 --- a/tests/e2e_tests/test_staking.py +++ b/tests/e2e_tests/test_staking.py @@ -3,6 +3,9 @@ from bittensor.core.chain_data.stake_info import StakeInfo from bittensor.utils.balance import Balance from tests.e2e_tests.utils.chain_interactions import ANY_BALANCE +from bittensor import logging + +logging.enable_info() def test_single_operation(subtensor, alice_wallet, bob_wallet): @@ -245,3 +248,275 @@ def test_batch_operations(subtensor, alice_wallet, bob_wallet): bob_wallet.coldkey.ss58_address: Balance.from_tao(999_998), } assert balances[alice_wallet.coldkey.ss58_address] > alice_balance + + +def test_safe_staking_scenarios(subtensor, alice_wallet, bob_wallet): + """ + Tests safe staking scenarios with different parameters. + + For both staking and unstaking: + 1. Fails with strict threshold (0.5%) and no partial staking + 2. Succeeds with strict threshold (0.5%) and partial staking allowed + 3. Succeeds with lenient threshold (10% and 30%) and no partial staking + """ + netuid = 2 + # Register root as Alice - the subnet owner and validator + assert subtensor.register_subnet(alice_wallet) + + # Verify subnet created successfully + assert subtensor.subnet_exists(netuid), "Subnet wasn't created successfully" + + subtensor.burned_register( + alice_wallet, + netuid=netuid, + wait_for_inclusion=True, + wait_for_finalization=True, + ) + subtensor.burned_register( + bob_wallet, + netuid=netuid, + wait_for_inclusion=True, + wait_for_finalization=True, + ) + + initial_stake = subtensor.get_stake( + alice_wallet.coldkey.ss58_address, + bob_wallet.hotkey.ss58_address, + netuid=netuid, + ) + assert initial_stake == Balance(0) + + # Test Staking Scenarios + stake_amount = Balance.from_tao(100) + + # 1. Strict params - should fail + success = subtensor.add_stake( + alice_wallet, + bob_wallet.hotkey.ss58_address, + netuid=netuid, + amount=stake_amount, + wait_for_inclusion=True, + wait_for_finalization=True, + safe_staking=True, + rate_threshold=0.005, # 0.5% + allow_partial_stake=False, + ) + assert success is False + + current_stake = subtensor.get_stake( + alice_wallet.coldkey.ss58_address, + bob_wallet.hotkey.ss58_address, + netuid=netuid, + ) + assert current_stake == Balance(0), "Stake should not change after failed attempt" + + # 2. Partial allowed - should succeed partially + success = subtensor.add_stake( + alice_wallet, + bob_wallet.hotkey.ss58_address, + netuid=netuid, + amount=stake_amount, + wait_for_inclusion=True, + wait_for_finalization=True, + safe_staking=True, + rate_threshold=0.005, # 0.5% + allow_partial_stake=True, + ) + assert success is True + + partial_stake = subtensor.get_stake( + alice_wallet.coldkey.ss58_address, + bob_wallet.hotkey.ss58_address, + netuid=netuid, + ) + assert partial_stake > Balance(0), "Partial stake should be added" + assert ( + partial_stake < stake_amount + ), "Partial stake should be less than requested amount" + + # 3. Higher threshold - should succeed fully + amount = Balance.from_tao(100) + success = subtensor.add_stake( + alice_wallet, + bob_wallet.hotkey.ss58_address, + netuid=netuid, + amount=amount, + wait_for_inclusion=True, + wait_for_finalization=True, + safe_staking=True, + rate_threshold=0.1, # 10% + allow_partial_stake=False, + ) + assert success is True + + full_stake = subtensor.get_stake( + alice_wallet.coldkey.ss58_address, + bob_wallet.hotkey.ss58_address, + netuid=netuid, + ) + + # Test Unstaking Scenarios + # 1. Strict params - should fail + success = subtensor.unstake( + alice_wallet, + bob_wallet.hotkey.ss58_address, + netuid=netuid, + amount=stake_amount, + wait_for_inclusion=True, + wait_for_finalization=True, + safe_staking=True, + rate_threshold=0.005, # 0.5% + allow_partial_stake=False, + ) + assert success is False + + current_stake = subtensor.get_stake( + alice_wallet.coldkey.ss58_address, + bob_wallet.hotkey.ss58_address, + netuid=netuid, + ) + assert ( + current_stake == full_stake + ), "Stake should not change after failed unstake attempt" + + # 2. Partial allowed - should succeed partially + success = subtensor.unstake( + alice_wallet, + bob_wallet.hotkey.ss58_address, + netuid=netuid, + amount=current_stake, + wait_for_inclusion=True, + wait_for_finalization=True, + safe_staking=True, + rate_threshold=0.005, # 0.5% + allow_partial_stake=True, + ) + assert success is True + + partial_unstake = subtensor.get_stake( + alice_wallet.coldkey.ss58_address, + bob_wallet.hotkey.ss58_address, + netuid=netuid, + ) + assert partial_unstake > Balance(0), "Some stake should remain" + + # 3. Higher threshold - should succeed fully + success = subtensor.unstake( + alice_wallet, + bob_wallet.hotkey.ss58_address, + netuid=netuid, + amount=partial_unstake, + wait_for_inclusion=True, + wait_for_finalization=True, + safe_staking=True, + rate_threshold=0.3, # 30% + allow_partial_stake=False, + ) + assert success is True + + +def test_safe_swap_stake_scenarios(subtensor, alice_wallet, bob_wallet): + """ + Tests safe swap stake scenarios with different parameters. + + Tests: + 1. Fails with strict threshold (0.5%) + 2. Succeeds with lenient threshold (10%) + """ + # Create new subnet (netuid 2) and register Alice + origin_netuid = 2 + assert subtensor.register_subnet(bob_wallet) + assert subtensor.subnet_exists(origin_netuid), "Subnet wasn't created successfully" + dest_netuid = 3 + assert subtensor.register_subnet(bob_wallet) + assert subtensor.subnet_exists(dest_netuid), "Subnet wasn't created successfully" + + # Register Alice on both subnets + subtensor.burned_register( + alice_wallet, + netuid=origin_netuid, + wait_for_inclusion=True, + wait_for_finalization=True, + ) + subtensor.burned_register( + alice_wallet, + netuid=dest_netuid, + wait_for_inclusion=True, + wait_for_finalization=True, + ) + + # Add initial stake to swap from + initial_stake_amount = Balance.from_tao(10_000) + success = subtensor.add_stake( + alice_wallet, + alice_wallet.hotkey.ss58_address, + netuid=origin_netuid, + amount=initial_stake_amount, + wait_for_inclusion=True, + wait_for_finalization=True, + ) + assert success is True + + origin_stake = subtensor.get_stake( + alice_wallet.coldkey.ss58_address, + alice_wallet.hotkey.ss58_address, + netuid=origin_netuid, + ) + assert origin_stake > Balance(0), "Origin stake should be non-zero" + + stake_swap_amount = Balance.from_tao(10_000) + # 1. Try swap with strict threshold and big amount- should fail + success = subtensor.swap_stake( + wallet=alice_wallet, + hotkey_ss58=alice_wallet.hotkey.ss58_address, + origin_netuid=origin_netuid, + destination_netuid=dest_netuid, + amount=stake_swap_amount, + wait_for_inclusion=True, + wait_for_finalization=True, + safe_staking=True, + rate_threshold=0.005, # 0.5% + allow_partial_stake=False, + ) + assert success is False + + # Verify no stake was moved + dest_stake = subtensor.get_stake( + alice_wallet.coldkey.ss58_address, + alice_wallet.hotkey.ss58_address, + netuid=dest_netuid, + ) + assert dest_stake == Balance( + 0 + ), "Destination stake should remain 0 after failed swap" + + # 2. Try swap with higher threshold and less amount - should succeed + stake_swap_amount = Balance.from_tao(100) + success = subtensor.swap_stake( + wallet=alice_wallet, + hotkey_ss58=alice_wallet.hotkey.ss58_address, + origin_netuid=origin_netuid, + destination_netuid=dest_netuid, + amount=stake_swap_amount, + wait_for_inclusion=True, + wait_for_finalization=True, + safe_staking=True, + rate_threshold=0.3, # 30% + allow_partial_stake=True, + ) + assert success is True + + # Verify stake was moved + origin_stake = subtensor.get_stake( + alice_wallet.coldkey.ss58_address, + alice_wallet.hotkey.ss58_address, + netuid=origin_netuid, + ) + dest_stake = subtensor.get_stake( + alice_wallet.coldkey.ss58_address, + alice_wallet.hotkey.ss58_address, + netuid=dest_netuid, + ) + assert dest_stake > Balance( + 0 + ), "Destination stake should be non-zero after successful swap" diff --git a/tests/unit_tests/test_subtensor.py b/tests/unit_tests/test_subtensor.py index cc47344eeb..1dce2d893b 100644 --- a/tests/unit_tests/test_subtensor.py +++ b/tests/unit_tests/test_subtensor.py @@ -2,14 +2,14 @@ # Copyright © 2024 Opentensor Foundation # # Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated -# documentation files (the “Software”), to deal in the Software without restriction, including without limitation +# documentation files (the "Software"), to deal in the Software without restriction, including without limitation # the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, # and to permit persons to whom the Software is furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in all copies or substantial portions of # the Software. # -# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO # THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL # THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER @@ -2854,6 +2854,9 @@ def test_add_stake_success(mocker, fake_wallet, subtensor): amount=fake_amount, wait_for_inclusion=True, wait_for_finalization=False, + safe_staking=False, + allow_partial_stake=False, + rate_threshold=0.005, ) # Assertions @@ -2865,6 +2868,48 @@ def test_add_stake_success(mocker, fake_wallet, subtensor): amount=Balance.from_rao(fake_amount), wait_for_inclusion=True, wait_for_finalization=False, + safe_staking=False, + allow_partial_stake=False, + rate_threshold=0.005, + ) + assert result == mock_add_stake_extrinsic.return_value + + +def test_add_stake_with_safe_staking(mocker, fake_wallet, subtensor): + """Test add_stake with safe staking parameters enabled.""" + # Prep + fake_hotkey_ss58 = "fake_hotkey" + fake_amount = 10.0 + fake_rate_threshold = 0.01 # 1% threshold + + mock_add_stake_extrinsic = mocker.patch.object( + subtensor_module, "add_stake_extrinsic" + ) + + # Call + result = subtensor.add_stake( + wallet=fake_wallet, + hotkey_ss58=fake_hotkey_ss58, + amount=fake_amount, + wait_for_inclusion=True, + wait_for_finalization=False, + safe_staking=True, + allow_partial_stake=False, + rate_threshold=fake_rate_threshold, + ) + + # Assertions + mock_add_stake_extrinsic.assert_called_once_with( + subtensor=subtensor, + wallet=fake_wallet, + hotkey_ss58=fake_hotkey_ss58, + netuid=None, + amount=Balance.from_rao(fake_amount), + wait_for_inclusion=True, + wait_for_finalization=False, + safe_staking=True, + allow_partial_stake=False, + rate_threshold=fake_rate_threshold, ) assert result == mock_add_stake_extrinsic.return_value @@ -2917,6 +2962,9 @@ def test_unstake_success(mocker, subtensor, fake_wallet): amount=fake_amount, wait_for_inclusion=True, wait_for_finalization=False, + safe_staking=False, + allow_partial_stake=False, + rate_threshold=0.005, ) # Assertions @@ -2928,10 +2976,136 @@ def test_unstake_success(mocker, subtensor, fake_wallet): amount=Balance.from_rao(fake_amount), wait_for_inclusion=True, wait_for_finalization=False, + safe_staking=False, + allow_partial_stake=False, + rate_threshold=0.005, ) assert result == mock_unstake_extrinsic.return_value +def test_unstake_with_safe_staking(mocker, subtensor, fake_wallet): + """Test unstake with safe staking parameters enabled.""" + fake_hotkey_ss58 = "hotkey_1" + fake_amount = 10.0 + fake_rate_threshold = 0.01 # 1% threshold + + mock_unstake_extrinsic = mocker.patch.object(subtensor_module, "unstake_extrinsic") + + # Call + result = subtensor.unstake( + wallet=fake_wallet, + hotkey_ss58=fake_hotkey_ss58, + amount=fake_amount, + wait_for_inclusion=True, + wait_for_finalization=False, + safe_staking=True, + allow_partial_stake=True, + rate_threshold=fake_rate_threshold, + ) + + # Assertions + mock_unstake_extrinsic.assert_called_once_with( + subtensor=subtensor, + wallet=fake_wallet, + hotkey_ss58=fake_hotkey_ss58, + netuid=None, + amount=Balance.from_rao(fake_amount), + wait_for_inclusion=True, + wait_for_finalization=False, + safe_staking=True, + allow_partial_stake=True, + rate_threshold=fake_rate_threshold, + ) + assert result == mock_unstake_extrinsic.return_value + + +def test_swap_stake_success(mocker, subtensor, fake_wallet): + """Test swap_stake operation is successful.""" + # Preps + fake_hotkey_ss58 = "hotkey_1" + fake_origin_netuid = 1 + fake_destination_netuid = 2 + fake_amount = 10.0 + + mock_swap_stake_extrinsic = mocker.patch.object( + subtensor_module, "swap_stake_extrinsic" + ) + + # Call + result = subtensor.swap_stake( + wallet=fake_wallet, + hotkey_ss58=fake_hotkey_ss58, + origin_netuid=fake_origin_netuid, + destination_netuid=fake_destination_netuid, + amount=fake_amount, + wait_for_inclusion=True, + wait_for_finalization=False, + safe_staking=False, + allow_partial_stake=False, + rate_threshold=0.005, + ) + + # Assertions + mock_swap_stake_extrinsic.assert_called_once_with( + subtensor=subtensor, + wallet=fake_wallet, + hotkey_ss58=fake_hotkey_ss58, + origin_netuid=fake_origin_netuid, + destination_netuid=fake_destination_netuid, + amount=Balance.from_rao(fake_amount), + wait_for_inclusion=True, + wait_for_finalization=False, + safe_staking=False, + allow_partial_stake=False, + rate_threshold=0.005, + ) + assert result == mock_swap_stake_extrinsic.return_value + + +def test_swap_stake_with_safe_staking(mocker, subtensor, fake_wallet): + """Test swap_stake with safe staking parameters enabled.""" + # Preps + fake_hotkey_ss58 = "hotkey_1" + fake_origin_netuid = 1 + fake_destination_netuid = 2 + fake_amount = 10.0 + fake_rate_threshold = 0.01 # 1% threshold + + mock_swap_stake_extrinsic = mocker.patch.object( + subtensor_module, "swap_stake_extrinsic" + ) + + # Call + result = subtensor.swap_stake( + wallet=fake_wallet, + hotkey_ss58=fake_hotkey_ss58, + origin_netuid=fake_origin_netuid, + destination_netuid=fake_destination_netuid, + amount=fake_amount, + wait_for_inclusion=True, + wait_for_finalization=False, + safe_staking=True, + allow_partial_stake=True, + rate_threshold=fake_rate_threshold, + ) + + # Assertions + mock_swap_stake_extrinsic.assert_called_once_with( + subtensor=subtensor, + wallet=fake_wallet, + hotkey_ss58=fake_hotkey_ss58, + origin_netuid=fake_origin_netuid, + destination_netuid=fake_destination_netuid, + amount=Balance.from_rao(fake_amount), + wait_for_inclusion=True, + wait_for_finalization=False, + safe_staking=True, + allow_partial_stake=True, + rate_threshold=fake_rate_threshold, + ) + assert result == mock_swap_stake_extrinsic.return_value + + def test_unstake_multiple_success(mocker, subtensor, fake_wallet): """Test unstake_multiple succeeds for all hotkeys.""" # Preps