diff --git a/bittensor/core/async_subtensor.py b/bittensor/core/async_subtensor.py index d1f61533a0..1588f2f236 100644 --- a/bittensor/core/async_subtensor.py +++ b/bittensor/core/async_subtensor.py @@ -520,8 +520,7 @@ async def get_transfer_fee( ) payment_info = {"partialFee": int(2e7)} # assume 0.02 Tao - fee = Balance.from_rao(payment_info["partialFee"]) - return fee + return Balance.from_rao(payment_info["partialFee"]) else: fee = Balance.from_rao(int(2e7)) logging.error( diff --git a/bittensor/utils/async_substrate_interface.py b/bittensor/utils/async_substrate_interface.py index 46fcc9347e..c3af691952 100644 --- a/bittensor/utils/async_substrate_interface.py +++ b/bittensor/utils/async_substrate_interface.py @@ -6,6 +6,7 @@ from hashlib import blake2b from typing import Optional, Any, Union, Callable, Awaitable, cast +import websockets from async_property import async_property from bittensor_wallet import Keypair from bt_decode import PortableRegistry, decode as decode_by_type_string, MetadataV15 @@ -20,7 +21,6 @@ BlockNotFound, ) from substrateinterface.storage import StorageKey -import websockets ResultHandler = Callable[[dict, Any], Awaitable[tuple[dict, bool]]] diff --git a/tests/unit_tests/test_async_subtensor.py b/tests/unit_tests/test_async_subtensor.py index 6b58684f70..dac1607a36 100644 --- a/tests/unit_tests/test_async_subtensor.py +++ b/tests/unit_tests/test_async_subtensor.py @@ -1,11 +1,9 @@ -from pickle import FALSE - import pytest from bittensor.core import async_subtensor -@pytest.fixture +@pytest.fixture(autouse=True) def subtensor(mocker): fake_async_substrate = mocker.AsyncMock( autospec=async_subtensor.AsyncSubstrateInterface @@ -43,6 +41,37 @@ def test_decode_ss58_tuples_in_proposal_vote_data(mocker): ] +def test__str__return(subtensor): + """Simply tests the result if printing subtensor instance.""" + # Asserts + assert ( + str(subtensor) + == "Network: finney, Chain: wss://entrypoint-finney.opentensor.ai:443" + ) + + +@pytest.mark.asyncio +async def test_async_subtensor_magic_methods(mocker): + """Tests async magic methods of AsyncSubtensor class.""" + # Preps + fake_async_substrate = mocker.AsyncMock( + autospec=async_subtensor.AsyncSubstrateInterface + ) + mocker.patch.object( + async_subtensor, "AsyncSubstrateInterface", return_value=fake_async_substrate + ) + + # Call + subtensor = async_subtensor.AsyncSubtensor(network="local") + async with subtensor: + pass + + # Asserts + fake_async_substrate.__aenter__.assert_called_once() + fake_async_substrate.__aexit__.assert_called_once() + fake_async_substrate.close.assert_awaited_once() + + @pytest.mark.asyncio async def test_encode_params(subtensor, mocker): """Tests encode_params happy path.""" @@ -291,3 +320,312 @@ async def test_get_delegates(subtensor, mocker, fake_hex_bytes_result, response) block_hash=None, reuse_block=True, ) + + +@pytest.mark.parametrize( + "fake_hex_bytes_result, response", [(None, []), ("zz001122", b"\xaa\xbb\xcc\xdd")] +) +@pytest.mark.asyncio +async def test_get_stake_info_for_coldkey( + subtensor, mocker, fake_hex_bytes_result, response +): + """Tests get_stake_info_for_coldkey method.""" + # Preps + fake_coldkey_ss58 = "fake_coldkey_58" + + mocked_ss58_to_vec_u8 = mocker.Mock() + async_subtensor.ss58_to_vec_u8 = mocked_ss58_to_vec_u8 + + mocked_query_runtime_api = mocker.AsyncMock( + autospec=subtensor.query_runtime_api, return_value=fake_hex_bytes_result + ) + subtensor.query_runtime_api = mocked_query_runtime_api + + mocked_stake_info_list_from_vec_u8 = mocker.Mock() + async_subtensor.StakeInfo.list_from_vec_u8 = mocked_stake_info_list_from_vec_u8 + + # Call + result = await subtensor.get_stake_info_for_coldkey( + coldkey_ss58=fake_coldkey_ss58, block_hash=None, reuse_block=True + ) + + # Asserts + if fake_hex_bytes_result: + assert result == mocked_stake_info_list_from_vec_u8.return_value + mocked_stake_info_list_from_vec_u8.assert_called_once_with( + bytes.fromhex(fake_hex_bytes_result[2:]) + ) + else: + assert result == response + + mocked_ss58_to_vec_u8.assert_called_once_with(fake_coldkey_ss58) + mocked_query_runtime_api.assert_called_once_with( + runtime_api="StakeInfoRuntimeApi", + method="get_stake_info_for_coldkey", + params=[mocked_ss58_to_vec_u8.return_value], + block_hash=None, + reuse_block=True, + ) + + +@pytest.mark.asyncio +async def test_get_stake_for_coldkey_and_hotkey(subtensor, mocker): + """Tests get_stake_for_coldkey_and_hotkey method.""" + # Preps + mocked_substrate_query = mocker.AsyncMock( + autospec=async_subtensor.AsyncSubstrateInterface.query + ) + subtensor.substrate.query = mocked_substrate_query + + spy_balance = mocker.spy(async_subtensor, "Balance") + + # Call + result = await subtensor.get_stake_for_coldkey_and_hotkey( + hotkey_ss58="hotkey", coldkey_ss58="coldkey", block_hash=None + ) + + # Asserts + mocked_substrate_query.assert_called_once_with( + module="SubtensorModule", + storage_function="Stake", + params=["hotkey", "coldkey"], + block_hash=None, + ) + assert result == spy_balance.from_rao.return_value + spy_balance.from_rao.assert_called_once_with(mocked_substrate_query.return_value) + + +@pytest.mark.asyncio +async def test_query_runtime_api(subtensor, mocker): + """Tests query_runtime_api method.""" + # Preps + fake_runtime_api = "DelegateInfoRuntimeApi" + fake_method = "get_delegated" + fake_params = [1, 2, 3] + fake_block_hash = None + reuse_block = False + + mocked_encode_params = mocker.AsyncMock() + subtensor.encode_params = mocked_encode_params + + mocked_rpc_request = mocker.AsyncMock( + autospec=async_subtensor.AsyncSubstrateInterface.rpc_request + ) + subtensor.substrate.rpc_request = mocked_rpc_request + + mocked_scalecodec = mocker.Mock(autospec=async_subtensor.scalecodec.ScaleBytes) + async_subtensor.scalecodec.ScaleBytes = mocked_scalecodec + + mocked_runtime_configuration = mocker.Mock( + autospec=async_subtensor.RuntimeConfiguration + ) + async_subtensor.RuntimeConfiguration = mocked_runtime_configuration + + mocked_load_type_registry_preset = mocker.Mock() + async_subtensor.load_type_registry_preset = mocked_load_type_registry_preset + + # Call + result = await subtensor.query_runtime_api( + runtime_api=fake_runtime_api, + method=fake_method, + params=fake_params, + block_hash=fake_block_hash, + reuse_block=reuse_block, + ) + + # Asserts + + mocked_encode_params.assert_called_once_with( + call_definition={ + "params": [{"name": "coldkey", "type": "Vec"}], + "type": "Vec", + }, + params=[1, 2, 3], + ) + mocked_rpc_request.assert_called_once_with( + method="state_call", + params=[f"{fake_runtime_api}_{fake_method}", mocked_encode_params.return_value], + reuse_block_hash=reuse_block, + ) + mocked_runtime_configuration.assert_called_once() + assert ( + mocked_runtime_configuration.return_value.update_type_registry.call_count == 2 + ) + + mocked_runtime_configuration.return_value.create_scale_object.assert_called_once_with( + "Vec", mocked_scalecodec.return_value + ) + + assert ( + result + == mocked_runtime_configuration.return_value.create_scale_object.return_value.decode.return_value + ) + + +@pytest.mark.asyncio +async def test_get_balance(subtensor, mocker): + """Tests get_balance method.""" + # Preps + fake_addresses = ("a1", "a2") + fake_block_hash = None + + mocked_substrate_create_storage_key = mocker.AsyncMock() + subtensor.substrate.create_storage_key = mocked_substrate_create_storage_key + + mocked_batch_0_call = mocker.Mock( + params=[ + 0, + ] + ) + mocked_batch_1_call = {"data": {"free": 1000}} + mocked_substrate_query_multi = mocker.AsyncMock( + return_value=[ + (mocked_batch_0_call, mocked_batch_1_call), + ] + ) + + subtensor.substrate.query_multi = mocked_substrate_query_multi + + # Call + result = await subtensor.get_balance(*fake_addresses, block_hash=fake_block_hash) + + assert mocked_substrate_create_storage_key.call_count == len(fake_addresses) + mocked_substrate_query_multi.assert_called_once() + assert result == {0: async_subtensor.Balance(1000)} + + +@pytest.mark.parametrize("balance", [100, 100.1]) +@pytest.mark.asyncio +async def test_get_transfer_fee(subtensor, mocker, balance): + """Tests get_transfer_fee method.""" + # Preps + fake_wallet = mocker.Mock(coldkeypub="coldkeypub", autospec=async_subtensor.Wallet) + fake_dest = "fake_dest" + fake_value = balance + + mocked_compose_call = mocker.AsyncMock() + subtensor.substrate.compose_call = mocked_compose_call + + mocked_get_payment_info = mocker.AsyncMock(return_value={"partialFee": 100}) + subtensor.substrate.get_payment_info = mocked_get_payment_info + + # Call + result = await subtensor.get_transfer_fee( + wallet=fake_wallet, dest=fake_dest, value=fake_value + ) + + # Assertions + mocked_compose_call.assert_awaited_once() + mocked_compose_call.assert_called_once_with( + call_module="Balances", + call_function="transfer_allow_death", + call_params={ + "dest": fake_dest, + "value": async_subtensor.Balance.from_rao(fake_value), + }, + ) + + assert isinstance(result, async_subtensor.Balance) + mocked_get_payment_info.assert_awaited_once() + mocked_get_payment_info.assert_called_once_with( + call=mocked_compose_call.return_value, keypair="coldkeypub" + ) + + +@pytest.mark.asyncio +async def test_get_transfer_fee_with_non_balance_accepted_value_type(subtensor, mocker): + """Tests get_transfer_fee method with non balance accepted value type.""" + # Preps + fake_wallet = mocker.Mock(coldkeypub="coldkeypub", autospec=async_subtensor.Wallet) + fake_dest = "fake_dest" + fake_value = "1000" + + # Call + result = await subtensor.get_transfer_fee( + wallet=fake_wallet, dest=fake_dest, value=fake_value + ) + + # Assertions + assert result == async_subtensor.Balance.from_rao(int(2e7)) + + +@pytest.mark.asyncio +async def test_get_transfer_with_exception(subtensor, mocker): + """Tests get_transfer_fee method handle Exception properly.""" + # Preps + fake_value = 123 + + mocked_compose_call = mocker.AsyncMock() + subtensor.substrate.compose_call = mocked_compose_call + subtensor.substrate.get_payment_info.side_effect = Exception + + # Call + result = await subtensor.get_transfer_fee( + wallet=mocker.Mock(), dest=mocker.Mock(), value=fake_value + ) + + # Assertions + assert result == async_subtensor.Balance.from_rao(int(2e7)) + + +@pytest.mark.asyncio +async def test_get_total_stake_for_coldkey(subtensor, mocker): + """Tests get_total_stake_for_coldkey method.""" + # Preps + fake_addresses = ("a1", "a2") + fake_block_hash = None + + mocked_substrate_create_storage_key = mocker.AsyncMock() + subtensor.substrate.create_storage_key = mocked_substrate_create_storage_key + + mocked_batch_0_call = mocker.Mock( + params=[ + 0, + ] + ) + mocked_batch_1_call = 0 + mocked_substrate_query_multi = mocker.AsyncMock( + return_value=[ + (mocked_batch_0_call, mocked_batch_1_call), + ] + ) + + subtensor.substrate.query_multi = mocked_substrate_query_multi + + # Call + result = await subtensor.get_total_stake_for_coldkey( + *fake_addresses, block_hash=fake_block_hash + ) + + assert mocked_substrate_create_storage_key.call_count == len(fake_addresses) + mocked_substrate_query_multi.assert_called_once() + assert result == {0: async_subtensor.Balance(mocked_batch_1_call)} + + +@pytest.mark.asyncio +async def test_get_total_stake_for_hotkey(subtensor, mocker): + """Tests get_total_stake_for_hotkey method.""" + # Preps + fake_addresses = ("a1", "a2") + fake_block_hash = None + reuse_block = True + + mocked_substrate_query_multiple = mocker.AsyncMock(return_value={0: 1}) + + subtensor.substrate.query_multiple = mocked_substrate_query_multiple + + # Call + result = await subtensor.get_total_stake_for_hotkey( + *fake_addresses, block_hash=fake_block_hash, reuse_block=reuse_block + ) + + # Assertions + mocked_substrate_query_multiple.assert_called_once_with( + params=list(fake_addresses), + module="SubtensorModule", + storage_function="TotalHotkeyStake", + block_hash=fake_block_hash, + reuse_block_hash=reuse_block, + ) + mocked_substrate_query_multiple.assert_called_once() + assert result == {0: async_subtensor.Balance(1)}