From 70454755b61b6ecc1fab8520741e74300af76f27 Mon Sep 17 00:00:00 2001 From: Benjamin Himes Date: Thu, 20 Feb 2025 14:59:16 +0200 Subject: [PATCH 1/4] Refactored the logic of `__aenter__` for `AsyncSubtensor` so that it doesn't close the websocket immediately. --- bittensor/core/async_subtensor.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bittensor/core/async_subtensor.py b/bittensor/core/async_subtensor.py index d83c373140..2f718d7c29 100644 --- 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." From c5b0986e36bb517ea491c24acf6b397887902ca2 Mon Sep 17 00:00:00 2001 From: Benjamin Himes Date: Thu, 20 Feb 2025 15:07:04 +0200 Subject: [PATCH 2/4] Fix tests. --- tests/unit_tests/test_async_subtensor.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/unit_tests/test_async_subtensor.py b/tests/unit_tests/test_async_subtensor.py index b328044dbc..ca2b6d187d 100644 --- 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( From 80f917425f0a662fdfa19679e375b2e8459b6e4f Mon Sep 17 00:00:00 2001 From: Benjamin Himes Date: Thu, 20 Feb 2025 20:03:50 +0200 Subject: [PATCH 3/4] Fix tests. --- tests/unit_tests/test_async_subtensor.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/test_async_subtensor.py b/tests/unit_tests/test_async_subtensor.py index ca2b6d187d..b6102b1cc9 100644 --- a/tests/unit_tests/test_async_subtensor.py +++ b/tests/unit_tests/test_async_subtensor.py @@ -144,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 @@ -157,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_called_once() @pytest.mark.asyncio From f7a93150003f042aa14f1a27429b34d64d4f2f12 Mon Sep 17 00:00:00 2001 From: Benjamin Himes Date: Thu, 20 Feb 2025 20:10:31 +0200 Subject: [PATCH 4/4] called -> awaited --- tests/unit_tests/test_async_subtensor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/test_async_subtensor.py b/tests/unit_tests/test_async_subtensor.py index b6102b1cc9..5123ce7607 100644 --- a/tests/unit_tests/test_async_subtensor.py +++ b/tests/unit_tests/test_async_subtensor.py @@ -157,7 +157,7 @@ async def test_async_subtensor_aenter_connection_refused_error( pass # Asserts - fake_async_substrate.initialize.assert_called_once() + fake_async_substrate.initialize.assert_awaited_once() @pytest.mark.asyncio