add unit tests

Author: Roman

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

tests/unit_tests/test_subtensor.py

@@ -3271,96 +3271,6 @@ def test_get_timestamp(mocker, subtensor):
     assert expected_result == actual_result
 
 
-def test_stake_fee_methods(mocker, subtensor):
-    """Test the three stake fee calculation methods."""
-    # Mock data
-    fake_hotkey = "hk1"
-    fake_coldkey = "ck1"
-    fake_amount = Balance.from_tao(100)
-    netuid = 1
-    fake_fee = 1_000_000
-
-    # Mock return fee
-    mock_query = mocker.patch.object(
-        subtensor,
-        "query_runtime_api",
-        side_effect=lambda runtime_api, method, params, block: (
-            fake_fee
-            if runtime_api == "StakeInfoRuntimeApi" and method == "get_stake_fee"
-            else None
-        ),
-    )
-
-    # get_stake_add_fee
-    result = subtensor.get_stake_add_fee(
-        amount=fake_amount,
-        netuid=netuid,
-        coldkey_ss58=fake_coldkey,
-        hotkey_ss58=fake_hotkey,
-    )
-    assert isinstance(result, Balance)
-    assert result == Balance.from_rao(fake_fee)
-    mock_query.assert_called_with(
-        runtime_api="StakeInfoRuntimeApi",
-        method="get_stake_fee",
-        params=[
-            None,
-            fake_coldkey,
-            (fake_hotkey, netuid),
-            fake_coldkey,
-            fake_amount.rao,
-        ],
-        block=None,
-    )
-
-    # get_unstake_fee
-    result = subtensor.get_unstake_fee(
-        amount=fake_amount,
-        netuid=netuid,
-        coldkey_ss58=fake_coldkey,
-        hotkey_ss58=fake_hotkey,
-    )
-    assert isinstance(result, Balance)
-    assert result == Balance.from_rao(fake_fee)
-    mock_query.assert_called_with(
-        runtime_api="StakeInfoRuntimeApi",
-        method="get_stake_fee",
-        params=[
-            (fake_hotkey, netuid),
-            fake_coldkey,
-            None,
-            fake_coldkey,
-            fake_amount.rao,
-        ],
-        block=None,
-    )
-
-    # get_stake_movement_fee
-    result = subtensor.get_stake_movement_fee(
-        amount=fake_amount,
-        origin_netuid=2,
-        origin_hotkey_ss58=fake_hotkey,
-        origin_coldkey_ss58=fake_coldkey,
-        destination_netuid=3,
-        destination_hotkey_ss58="hk2",
-        destination_coldkey_ss58=fake_coldkey,
-    )
-    assert isinstance(result, Balance)
-    assert result == Balance.from_rao(fake_fee)
-    mock_query.assert_called_with(
-        runtime_api="StakeInfoRuntimeApi",
-        method="get_stake_fee",
-        params=[
-            (fake_hotkey, 2),
-            fake_coldkey,
-            ("hk2", 3),
-            fake_coldkey,
-            fake_amount.rao,
-        ],
-        block=None,
-    )
-
-
 def test_get_owned_hotkeys_happy_path(subtensor, mocker):
     """Tests that the output of get_owned_hotkeys."""
     # Prep
@@ -4266,3 +4176,102 @@ def test_subnet(subtensor, mocker):
         {"netuid": netuid, "price": Balance.from_tao(100.0)}
     )
     assert result == mocked_di_from_dict.return_value
+
+
+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 = subtensor.get_stake_operations_fee(amount=amount, netuid=netuid)
+
+    # Assert
+    mocked_determine_block_hash.assert_called_once_with(block=None)
+    mocked_query_map.assert_called_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)
+
+
+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 = subtensor.get_stake_add_fee(
+        amount=amount,
+        netuid=netuid,
+        coldkey_ss58=mocker.Mock(),
+        hotkey_ss58=mocker.Mock(),
+    )
+
+    # Asserts
+    mocked_get_stake_operations_fee.assert_called_once_with(
+        amount=amount, netuid=netuid, block=None
+    )
+    assert result == mocked_get_stake_operations_fee.return_value
+
+
+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 = subtensor.get_unstake_fee(
+        amount=amount,
+        netuid=netuid,
+        coldkey_ss58=mocker.Mock(),
+        hotkey_ss58=mocker.Mock(),
+    )
+
+    # Asserts
+    mocked_get_stake_operations_fee.assert_called_once_with(
+        amount=amount, netuid=netuid, block=None
+    )
+    assert result == mocked_get_stake_operations_fee.return_value
+
+
+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 = 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_called_once_with(
+        amount=amount, netuid=netuid, block=None
+    )
+    assert result == mocked_get_stake_operations_fee.return_value