diff --git a/bittensor/core/async_subtensor.py b/bittensor/core/async_subtensor.py index e1f1d1b51d..e74bc92d22 100644 --- a/bittensor/core/async_subtensor.py +++ b/bittensor/core/async_subtensor.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import copy import ssl from datetime import datetime, timezone @@ -358,6 +359,19 @@ async def __aenter__(self): async def __aexit__(self, exc_type, exc_val, exc_tb): await self.substrate.close() + @contextlib.asynccontextmanager + async def _ensure_connection(self): + """Ensure the substrate connection is open for the duration of a single call + without closing it afterwards. + + Unlike ``async with self:``, this must be used by methods that run on a + long-lived, shared connection: ``__aexit__`` calls ``self.substrate.close()``, + so ``async with self:`` inside a method would close the caller's connection as + soon as the method returns, breaking every subsequent query on that subtensor. + """ + await self.initialize() + yield + # Helpers ========================================================================================================== async def _decode_crowdloan_entry( @@ -7216,7 +7230,7 @@ async def burned_register( Notes: - Rate Limits: """ - async with self: + async with self._ensure_connection(): if netuid == 0: return await root_register_extrinsic( subtensor=self, @@ -8507,7 +8521,7 @@ async def register( Notes: - Rate Limits: """ - async with self: + async with self._ensure_connection(): if netuid == 0: return await root_register_extrinsic( subtensor=self, @@ -8586,7 +8600,7 @@ async def register_limit( - Rate Limits: """ check_balance_amount(limit_price) - async with self: + async with self._ensure_connection(): return await register_limit_extrinsic( subtensor=self, wallet=wallet, diff --git a/mypy.ini b/mypy.ini index d38bdc7172..054851282e 100644 --- a/mypy.ini +++ b/mypy.ini @@ -1,6 +1,13 @@ [mypy] ignore_missing_imports = True ignore_errors = True +follow_imports_for_stubs = True + +[mypy-numpy.*] +follow_imports = skip + +[mypy-numpy] +follow_imports = skip [mypy-*.axon.*] ignore_errors = False diff --git a/tests/unit_tests/test_async_subtensor.py b/tests/unit_tests/test_async_subtensor.py index 94029b20ab..156042f565 100644 --- a/tests/unit_tests/test_async_subtensor.py +++ b/tests/unit_tests/test_async_subtensor.py @@ -194,6 +194,31 @@ async def test_burned_register(subtensor, fake_wallet, mocker): ) +@pytest.mark.asyncio +async def test_register_family_does_not_close_shared_connection( + subtensor, fake_wallet, mocker +): + """register/burned_register/register_limit must not close the substrate connection + they share with the caller. The old ``async with self:`` ran ``__aexit__`` -> + ``substrate.close()`` on return, breaking every later query on the same subtensor.""" + mocker.patch.object(subtensor, "compose_call") + mocker.patch.object( + subtensor, "sign_and_send_extrinsic", return_value=ExtrinsicResponse(True, "") + ) + mocker.patch.object( + subtensor, + "get_neuron_for_pubkey_and_subnet", + return_value=NeuronInfo.get_null_neuron(), + ) + mocker.patch.object(subtensor, "get_balance", return_value=Balance.from_tao(1)) + mocker.patch.object(subtensor, "recycle") + + await subtensor.burned_register(wallet=fake_wallet, netuid=14) + + # the connection the caller is sharing must be left open + subtensor.substrate.close.assert_not_called() + + @pytest.mark.asyncio async def test_burned_register_on_root(mock_substrate, subtensor, fake_wallet, mocker): mock_substrate.submit_extrinsic.return_value = mocker.AsyncMock(