Merge pull request #2684 from opentensor/feat/thewhaleking/refactor-a…

…enter-logic Refactor AsyncSubtensor aenter logic
opentensor · Feb 20, 2025 · 8aa8f4c · 8aa8f4c
2 parents 4bebb26 + f7a9315
commit 8aa8f4c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git a/bittensor/core/async_subtensor.py b/bittensor/core/async_subtensor.py
@@ -173,8 +173,8 @@ async def __aenter__(self):
             f"[magenta]Connecting to Substrate:[/magenta] [blue]{self}[/blue][magenta]...[/magenta]"
         )
         try:
-            async with self.substrate:
-                return self
+            await self.substrate.initialize()
+            return self
         except TimeoutError:
             logging.error(
                 f"[red]Error[/red]: Timeout occurred connecting to substrate."

diff --git a/tests/unit_tests/test_async_subtensor.py b/tests/unit_tests/test_async_subtensor.py
@@ -128,9 +128,8 @@ async def test_async_subtensor_magic_methods(mocker):
         pass
 
     # Asserts
-    fake_async_substrate.__aenter__.assert_called_once()
-    fake_async_substrate.__aexit__.assert_called_once()
-    fake_async_substrate.close.assert_awaited_once()
+    fake_async_substrate.initialize.assert_called_once()
+    fake_async_substrate.close.assert_called_once()
 
 
 @pytest.mark.parametrize(
@@ -145,7 +144,7 @@ async def test_async_subtensor_aenter_connection_refused_error(
     # Preps
     fake_async_substrate = mocker.AsyncMock(
         autospec=async_subtensor.AsyncSubstrateInterface,
-        __aenter__=mocker.AsyncMock(side_effect=error),
+        initialize=mocker.AsyncMock(side_effect=error),
     )
     mocker.patch.object(
         async_subtensor, "AsyncSubstrateInterface", return_value=fake_async_substrate
@@ -158,7 +157,7 @@ async def test_async_subtensor_aenter_connection_refused_error(
             pass
 
     # Asserts
-    fake_async_substrate.__aenter__.assert_called_once()
+    fake_async_substrate.initialize.assert_awaited_once()
 
 
 @pytest.mark.asyncio