Merge pull request #2684 from opentensor/feat/thewhaleking/refactor-aenter-logic

Refactor AsyncSubtensor aenter logic

Author: thewhaleking

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."

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