Stop register methods from closing the shared async subtensor connection

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: <https://docs.learnbittensor.org/learn/chain-rate-limits#registration-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: <https://docs.learnbittensor.org/learn/chain-rate-limits#registration-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: <https://docs.learnbittensor.org/learn/chain-rate-limits#registration-rate-limits>
         """
         check_balance_amount(limit_price)
-        async with self:
+        async with self._ensure_connection():
             return await register_limit_extrinsic(
                 subtensor=self,
                 wallet=wallet,

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

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(