Merge pull request #2961 from opentensor/feat/roman/update-fee-logic

Update staking fee logic

Author: basfroman

bittensor/core/async_subtensor.py

@@ -117,7 +117,11 @@
     price_to_tick,
     LiquidityPosition,
 )
-from bittensor.utils.weight_utils import generate_weight_hash, convert_uids_and_weights
+from bittensor.utils.weight_utils import (
+    generate_weight_hash,
+    convert_uids_and_weights,
+    U16_MAX,
+)
 
 if TYPE_CHECKING:
     from async_substrate_interface.types import ScaleObj
@@ -2522,6 +2526,7 @@ async def get_stake(
 
         return Balance.from_rao(int(stake)).set_unit(netuid=netuid)
 
+    # TODO: remove unused parameters in SDK.v10
     async def get_stake_add_fee(
         self,
         amount: Balance,
@@ -2543,19 +2548,9 @@ async def get_stake_add_fee(
         Returns:
             The calculated stake fee as a Balance object
         """
-        result = await self.query_runtime_api(
-            runtime_api="StakeInfoRuntimeApi",
-            method="get_stake_fee",
-            params=[
-                None,
-                coldkey_ss58,
-                (hotkey_ss58, netuid),
-                coldkey_ss58,
-                amount.rao,
-            ],
-            block=block,
+        return await self.get_stake_operations_fee(
+            amount=amount, netuid=netuid, block=block
         )
-        return Balance.from_rao(result)
 
     async def get_subnet_info(
         self,
@@ -2669,6 +2664,7 @@ async def get_subnet_prices(
         prices.update({0: Balance.from_tao(1)})
         return prices
 
+    # TODO: remove unused parameters in SDK.v10
     async def get_unstake_fee(
         self,
         amount: Balance,
@@ -2690,20 +2686,11 @@ async def get_unstake_fee(
         Returns:
             The calculated stake fee as a Balance object
         """
-        result = await self.query_runtime_api(
-            runtime_api="StakeInfoRuntimeApi",
-            method="get_stake_fee",
-            params=[
-                None,
-                coldkey_ss58,
-                (hotkey_ss58, netuid),
-                coldkey_ss58,
-                amount.rao,
-            ],
-            block=block,
+        return await self.get_stake_operations_fee(
+            amount=amount, netuid=netuid, block=block
         )
-        return Balance.from_rao(result)
 
+    # TODO: remove unused parameters in SDK.v10
     async def get_stake_movement_fee(
         self,
         amount: Balance,
@@ -2731,19 +2718,9 @@ async def get_stake_movement_fee(
         Returns:
             The calculated stake fee as a Balance object
         """
-        result = await self.query_runtime_api(
-            runtime_api="StakeInfoRuntimeApi",
-            method="get_stake_fee",
-            params=[
-                (origin_hotkey_ss58, origin_netuid),
-                origin_coldkey_ss58,
-                (destination_hotkey_ss58, destination_netuid),
-                destination_coldkey_ss58,
-                amount.rao,
-            ],
-            block=block,
+        return await self.get_stake_operations_fee(
+            amount=amount, netuid=origin_netuid, block=block
         )
-        return Balance.from_rao(result)
 
     async def get_stake_for_coldkey_and_hotkey(
         self,
@@ -2864,6 +2841,39 @@ async def get_stake_for_hotkey(
 
     get_hotkey_stake = get_stake_for_hotkey
 
+    async def get_stake_operations_fee(
+        self,
+        amount: Balance,
+        netuid: int,
+        block: Optional[int] = None,
+        block_hash: Optional[str] = None,
+        reuse_block: bool = False,
+    ):
+        """Returns fee for any stake operation in specified subnet.
+
+        Args:
+            amount: Amount of stake to add in Alpha/TAO.
+            netuid: Netuid of subnet.
+            block: The block number at which to query the stake information. Do not specify if also specifying
+                block_hash or reuse_block.
+            block_hash: The hash of the blockchain block number for the query. Do not specify if also specifying block
+                or reuse_block.
+            reuse_block: Whether to reuse for this query the last-used block. Do not specify if also specifying block
+                or block_hash.
+        Returns:
+            The calculated stake fee as a Balance object.
+        """
+        block_hash = await self.determine_block_hash(
+            block=block, block_hash=block_hash, reuse_block=reuse_block
+        )
+        result = await self.substrate.query(
+            module="Swap",
+            storage_function="FeeRate",
+            params=[netuid],
+            block_hash=block_hash,
+        )
+        return amount * (result.value / U16_MAX)
+
     async def get_subnet_burn_cost(
         self,
         block: Optional[int] = None,

bittensor/core/subtensor.py

@@ -120,7 +120,11 @@
     price_to_tick,
     LiquidityPosition,
 )
-from bittensor.utils.weight_utils import generate_weight_hash, convert_uids_and_weights
+from bittensor.utils.weight_utils import (
+    generate_weight_hash,
+    convert_uids_and_weights,
+    U16_MAX,
+)
 
 if TYPE_CHECKING:
     from bittensor_wallet import Wallet
@@ -1756,6 +1760,7 @@ def get_stake(
 
         return Balance.from_rao(int(stake)).set_unit(netuid=netuid)
 
+    # TODO: remove unused parameters in SDK.v10
     def get_stake_add_fee(
         self,
         amount: Balance,
@@ -1777,19 +1782,7 @@ def get_stake_add_fee(
         Returns:
             The calculated stake fee as a Balance object
         """
-        result = self.query_runtime_api(
-            runtime_api="StakeInfoRuntimeApi",
-            method="get_stake_fee",
-            params=[
-                None,
-                coldkey_ss58,
-                (hotkey_ss58, netuid),
-                coldkey_ss58,
-                amount.rao,
-            ],
-            block=block,
-        )
-        return Balance.from_rao(result)
+        return self.get_stake_operations_fee(amount=amount, netuid=netuid, block=block)
 
     def get_subnet_info(
         self, netuid: int, block: Optional[int] = None
@@ -1882,6 +1875,7 @@ def get_subnet_prices(
         prices.update({0: Balance.from_tao(1)})
         return prices
 
+    # TODO: remove unused parameters in SDK.v10
     def get_unstake_fee(
         self,
         amount: Balance,
@@ -1903,20 +1897,9 @@ def get_unstake_fee(
         Returns:
             The calculated stake fee as a Balance object
         """
-        result = self.query_runtime_api(
-            runtime_api="StakeInfoRuntimeApi",
-            method="get_stake_fee",
-            params=[
-                (hotkey_ss58, netuid),
-                coldkey_ss58,
-                None,
-                coldkey_ss58,
-                amount.rao,
-            ],
-            block=block,
-        )
-        return Balance.from_rao(result)
+        return self.get_stake_operations_fee(amount=amount, netuid=netuid, block=block)
 
+    # TODO: remove unused parameters in SDK.v10
     def get_stake_movement_fee(
         self,
         amount: Balance,
@@ -1944,19 +1927,9 @@ def get_stake_movement_fee(
         Returns:
             The calculated stake fee as a Balance object
         """
-        result = self.query_runtime_api(
-            runtime_api="StakeInfoRuntimeApi",
-            method="get_stake_fee",
-            params=[
-                (origin_hotkey_ss58, origin_netuid),
-                origin_coldkey_ss58,
-                (destination_hotkey_ss58, destination_netuid),
-                destination_coldkey_ss58,
-                amount.rao,
-            ],
-            block=block,
+        return self.get_stake_operations_fee(
+            amount=amount, netuid=origin_netuid, block=block
         )
-        return Balance.from_rao(result)
 
     def get_stake_for_coldkey_and_hotkey(
         self,
@@ -2044,6 +2017,31 @@ def get_stake_for_hotkey(
 
     get_hotkey_stake = get_stake_for_hotkey
 
+    def get_stake_operations_fee(
+        self,
+        amount: Balance,
+        netuid: int,
+        block: Optional[int] = None,
+    ):
+        """Returns fee for any stake operation in specified subnet.
+
+        Args:
+            amount: Amount of stake to add in Alpha/TAO.
+            netuid: Netuid of subnet.
+            block: Block number at which to perform the calculation.
+
+        Returns:
+            The calculated stake fee as a Balance object.
+        """
+        block_hash = self.determine_block_hash(block=block)
+        result = self.substrate.query(
+            module="Swap",
+            storage_function="FeeRate",
+            params=[netuid],
+            block_hash=block_hash,
+        )
+        return amount * (result.value / U16_MAX)
+
     def get_subnet_burn_cost(self, block: Optional[int] = None) -> Optional[Balance]:
         """
         Retrieves the burn cost for registering a new subnet within the Bittensor network. This cost represents the

bittensor/core/subtensor_api/staking.py

@@ -19,6 +19,7 @@ def __init__(self, subtensor: Union["_Subtensor", "_AsyncSubtensor"]):
         )
         self.get_stake_info_for_coldkey = subtensor.get_stake_info_for_coldkey
         self.get_stake_movement_fee = subtensor.get_stake_movement_fee
+        self.get_stake_operations_fee = subtensor.get_stake_operations_fee
         self.get_unstake_fee = subtensor.get_unstake_fee
         self.unstake = subtensor.unstake
         self.unstake_all = subtensor.unstake_all

bittensor/core/subtensor_api/utils.py

@@ -86,6 +86,7 @@ def add_legacy_methods(subtensor: "SubtensorApi"):
         subtensor._subtensor.get_stake_info_for_coldkey
     )
     subtensor.get_stake_movement_fee = subtensor._subtensor.get_stake_movement_fee
+    subtensor.get_stake_operations_fee = subtensor._subtensor.get_stake_operations_fee
     subtensor.get_subnet_burn_cost = subtensor._subtensor.get_subnet_burn_cost
     subtensor.get_subnet_hyperparameters = (
         subtensor._subtensor.get_subnet_hyperparameters

tests/e2e_tests/test_stake_fee.py

@@ -45,7 +45,9 @@ async def test_stake_fee_api(local_chain, subtensor, alice_wallet, bob_wallet):
         hotkey_ss58=bob_wallet.hotkey.ss58_address,
     )
     assert isinstance(unstake_fee_root, Balance), "Stake fee should be a Balance object"
-    assert unstake_fee_root == 0, "Root unstake fee should be 0."
+    assert unstake_fee_root == Balance.from_rao(299076829), (
+        "Root unstake fee should be 0."
+    )
 
     # Test various stake movement scenarios
     movement_scenarios = [

tests/unit_tests/test_async_subtensor.py

@@ -3955,3 +3955,108 @@ async def test_subnet(subtensor, mocker):
         {"netuid": netuid, "price": Balance.from_tao(100.0)}
     )
     assert result == mocked_di_from_dict.return_value
+
+
+@pytest.mark.asyncio
+async def test_get_stake_operations_fee(subtensor, mocker):
+    """Verify that `get_stake_operations_fee` calls proper methods and returns the correct value."""
+    # Preps
+    netuid = 1
+    amount = Balance.from_rao(100_000_000_000)  # 100 Tao
+    mocked_determine_block_hash = mocker.patch.object(subtensor, "determine_block_hash")
+    mocked_query_map = mocker.patch.object(
+        subtensor.substrate, "query", return_value=mocker.Mock(value=196)
+    )
+
+    # Call
+    result = await subtensor.get_stake_operations_fee(netuid=netuid, amount=amount)
+
+    # Assert
+    mocked_determine_block_hash.assert_awaited_once_with(
+        block=None, block_hash=None, reuse_block=False
+    )
+    mocked_query_map.assert_awaited_once_with(
+        module="Swap",
+        storage_function="FeeRate",
+        params=[netuid],
+        block_hash=mocked_determine_block_hash.return_value,
+    )
+    assert result == Balance.from_rao(299076829).set_unit(netuid)
+
+
+@pytest.mark.asyncio
+async def test_get_stake_add_fee(subtensor, mocker):
+    """Verify that `get_stake_add_fee` calls proper methods and returns the correct value."""
+    # Preps
+    netuid = mocker.Mock()
+    amount = mocker.Mock()
+    mocked_get_stake_operations_fee = mocker.patch.object(
+        subtensor, "get_stake_operations_fee"
+    )
+
+    # Call
+    result = await subtensor.get_stake_add_fee(
+        amount=amount,
+        netuid=netuid,
+        coldkey_ss58=mocker.Mock(),
+        hotkey_ss58=mocker.Mock(),
+    )
+
+    # Asserts
+    mocked_get_stake_operations_fee.assert_awaited_once_with(
+        amount=amount, netuid=netuid, block=None
+    )
+    assert result == mocked_get_stake_operations_fee.return_value
+
+
+@pytest.mark.asyncio
+async def test_get_unstake_fee(subtensor, mocker):
+    """Verify that `get_unstake_fee` calls proper methods and returns the correct value."""
+    # Preps
+    netuid = mocker.Mock()
+    amount = mocker.Mock()
+    mocked_get_stake_operations_fee = mocker.patch.object(
+        subtensor, "get_stake_operations_fee"
+    )
+
+    # Call
+    result = await subtensor.get_unstake_fee(
+        amount=amount,
+        netuid=netuid,
+        coldkey_ss58=mocker.Mock(),
+        hotkey_ss58=mocker.Mock(),
+    )
+
+    # Asserts
+    mocked_get_stake_operations_fee.assert_awaited_once_with(
+        amount=amount, netuid=netuid, block=None
+    )
+    assert result == mocked_get_stake_operations_fee.return_value
+
+
+@pytest.mark.asyncio
+async def test_get_stake_movement_fee(subtensor, mocker):
+    """Verify that `get_stake_movement_fee` calls proper methods and returns the correct value."""
+    # Preps
+    netuid = mocker.Mock()
+    amount = mocker.Mock()
+    mocked_get_stake_operations_fee = mocker.patch.object(
+        subtensor, "get_stake_operations_fee"
+    )
+
+    # Call
+    result = await subtensor.get_stake_movement_fee(
+        amount=amount,
+        origin_netuid=netuid,
+        origin_hotkey_ss58=mocker.Mock(),
+        origin_coldkey_ss58=mocker.Mock(),
+        destination_netuid=mocker.Mock(),
+        destination_hotkey_ss58=mocker.Mock(),
+        destination_coldkey_ss58=mocker.Mock(),
+    )
+
+    # Asserts
+    mocked_get_stake_operations_fee.assert_awaited_once_with(
+        amount=amount, netuid=netuid, block=None
+    )
+    assert result == mocked_get_stake_operations_fee.return_value