Tests

Author: thewhaleking

async_substrate_interface/async_substrate.py

@@ -607,7 +607,6 @@ async def _wait_with_activity_timeout(self, coro, timeout: float):
         """
         activity_task = asyncio.create_task(self._last_activity.wait())
 
-        # Handle both coroutines and tasks
         if isinstance(coro, asyncio.Task):
             main_task = coro
         else:
@@ -620,25 +619,22 @@ async def _wait_with_activity_timeout(self, coro, timeout: float):
                 return_when=asyncio.FIRST_COMPLETED,
             )
 
-            if not done:  # Timeout occurred
+            if not done:
                 logger.debug(f"Activity timeout after {timeout}s, no activity detected")
                 for task in pending:
                     task.cancel()
                 raise TimeoutError()
 
-            # Check which completed
             if main_task in done:
                 activity_task.cancel()
 
-                # Check if the task raised an exception
                 exc = main_task.exception()
                 if exc is not None:
                     raise exc
                 else:
                     return main_task.result()
-            else:  # activity_task completed (activity occurred elsewhere)
+            else:
                 logger.debug("Activity detected, resetting timeout")
-                # Recursively wait again with fresh timeout
                 return await self._wait_with_activity_timeout(main_task, timeout)
 
         except asyncio.CancelledError:
@@ -849,14 +845,12 @@ async def _start_receiving(self, ws: ClientConnection) -> Exception:
                     ws.recv(decode=False), self.retry_timeout
                 )
                 await self._reset_activity_timer()
-                # reset the counter once we successfully receive something back
                 self._attempts = 0
                 await self._recv(recd)
         except websockets.exceptions.ConnectionClosedOK as e:
             logger.debug("ConnectionClosedOK")
             return e
         except Exception as e:
-            logger.exception("Receiving exception", exc_info=e)
             if isinstance(e, ssl.SSLError):
                 e = ConnectionClosed
             if not isinstance(
@@ -890,7 +884,6 @@ async def _start_sending(self, ws) -> Exception:
                 logger.debug("Sent to websocket")
                 await self._reset_activity_timer()
         except Exception as e:
-            logger.exception("Maybe timeout? 769", exc_info=e)
             if isinstance(e, ssl.SSLError):
                 e = ConnectionClosed
             if not isinstance(

tests/helpers/proxy_server.py

@@ -1,7 +1,10 @@
-import asyncio
+import logging
+import time
 
-from websockets.asyncio.server import serve, ServerConnection
-from websockets.asyncio.client import connect
+from websockets.sync.server import serve, ServerConnection
+from websockets.sync.client import connect
+
+logger = logging.getLogger("websockets.proxy")
 
 
 class ProxyServer:
@@ -14,38 +17,41 @@ def __init__(self, upstream: str, time_til_pause: float, time_til_resume: float)
         self.shutdown_time = 0
         self.resume_time = 0
 
-    async def connect(self):
-        self.upstream_connection = await connect(self.upstream_server)
-        self.connection_time = asyncio.get_running_loop().time()
+    def connect(self):
+        self.upstream_connection = connect(self.upstream_server)
+        self.connection_time = time.time()
         self.shutdown_time = self.connection_time + self.time_til_pause
         self.resume_time = self.shutdown_time + self.time_til_resume
 
-    async def close(self):
+    def close(self):
         if self.upstream_connection:
-            await self.upstream_connection.close()
-
-    async def proxy_request(self, websocket: ServerConnection):
-        async for message in websocket:
-            print(websocket)
-            await self.upstream_connection.send(message)
-            recd = await self.upstream_connection.recv()
-            current_time = asyncio.get_running_loop().time()
+            self.upstream_connection.close()
+        self.server.shutdown()
+
+    def proxy_request(self, websocket: ServerConnection):
+        for message in websocket:
+            self.upstream_connection.send(message)
+            recd = self.upstream_connection.recv()
+            current_time = time.time()
             if self.shutdown_time < current_time < self.resume_time:
-                print("Pausing")
-                await asyncio.sleep(self.time_til_resume)
-            await websocket.send(recd)
-            # await websocket.send(message)
+                logger.info("Pausing")
+                time.sleep(self.time_til_resume)
+            websocket.send(recd)
+
+    def serve(self):
+        with serve(self.proxy_request, "localhost", 8080) as self.server:
+            self.server.serve_forever()
 
-    async def serve(self):
-        async with serve(self.proxy_request, "localhost", 8080) as server:
-            await server.serve_forever()
+    def connect_and_serve(self):
+        self.connect()
+        self.serve()
 
 
-async def main():
-    proxy = ProxyServer("wss://archive.sub.latent.to", 20, 30)
-    await proxy.connect()
-    await proxy.serve()
+def run_proxy_server(time_til_pause: float = 20.0, time_til_resume: float = 30.0):
+    proxy = ProxyServer("wss://archive.sub.latent.to", time_til_pause, time_til_resume)
+    proxy.connect()
+    proxy.serve()
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    run_proxy_server()

tests/integration_tests/test_async_substrate_interface.py

@@ -1,12 +1,16 @@
 import asyncio
+import logging
+import os.path
 import time
+import threading
 
 import pytest
 from scalecodec import ss58_encode
 
-from async_substrate_interface.async_substrate import AsyncSubstrateInterface
+from async_substrate_interface.async_substrate import AsyncSubstrateInterface, logger
 from async_substrate_interface.types import ScaleObj
 from tests.helpers.settings import ARCHIVE_ENTRYPOINT, LATENT_LITE_ENTRYPOINT
+from tests.helpers.proxy_server import ProxyServer
 
 
 @pytest.mark.asyncio
@@ -174,3 +178,120 @@ async def test_query_map_with_odd_number_of_params():
         first_record = qm.records[0]
         assert len(first_record) == 2
         assert len(first_record[0]) == 4
+
+
+@pytest.mark.asyncio
+async def test_improved_reconnection_():
+    ws_logger_path = "/tmp/websockets-proxy-test"
+    ws_logger = logging.getLogger("websockets.proxy")
+    if os.path.exists(ws_logger_path):
+        os.remove(ws_logger_path)
+    ws_logger.setLevel(logging.INFO)
+    ws_logger.addHandler(logging.FileHandler(ws_logger_path))
+
+    asi_logger_path = "/tmp/async-substrate-interface-test"
+    if os.path.exists(asi_logger_path):
+        os.remove(asi_logger_path)
+    logger.setLevel(logging.DEBUG)
+    logger.addHandler(logging.FileHandler(asi_logger_path))
+
+    proxy = ProxyServer("wss://archive.sub.latent.to", 10, 20)
+    proxy.connect()
+    server_thread = threading.Thread(target=proxy.serve)
+    server_thread.daemon = True
+    server_thread.start()
+
+    await asyncio.sleep(3)  # give the server start up time
+
+    try:
+        async with AsyncSubstrateInterface(
+            "ws://localhost:8080",
+            ss58_format=42,
+            chain_name="Bittensor",
+            retry_timeout=10.0,
+            ws_shutdown_timer=None,
+        ) as substrate:
+            blocks_to_check = [
+                5215000,
+                5215001,
+                5215002,
+                5215003,
+                5215004,
+                5215005,
+                5215006,
+            ]
+            tasks = []
+            for block in blocks_to_check:
+                block_hash = await substrate.get_block_hash(block_id=block)
+                tasks.append(
+                    substrate.query_map(
+                        "SubtensorModule", "TotalHotkeyShares", block_hash=block_hash
+                    )
+                )
+            records = await asyncio.gather(*tasks)
+            assert len(records) == len(blocks_to_check)
+            await substrate.close()
+            with open(ws_logger_path, "r") as f:
+                assert "Pausing" in f.read()
+            with open(asi_logger_path, "r") as f:
+                assert "Timeout/ConnectionClosed occurred." in f.read()
+    finally:
+        proxy.stop()
+        server_thread.join(timeout=5)
+
+
+@pytest.mark.asyncio
+async def test_improved_reconnection():
+    ws_logger_path = "/tmp/websockets-proxy-test"
+    ws_logger = logging.getLogger("websockets.proxy")
+    if os.path.exists(ws_logger_path):
+        os.remove(ws_logger_path)
+    ws_logger.setLevel(logging.INFO)
+    ws_logger.addHandler(logging.FileHandler(ws_logger_path))
+
+    asi_logger_path = "/tmp/async-substrate-interface-test"
+    if os.path.exists(asi_logger_path):
+        os.remove(asi_logger_path)
+    logger.setLevel(logging.DEBUG)
+    logger.addHandler(logging.FileHandler(asi_logger_path))
+
+    proxy = ProxyServer("wss://archive.sub.latent.to", 10, 20)
+
+    server_thread = threading.Thread(target=proxy.connect_and_serve)
+    server_thread.start()
+    await asyncio.sleep(3)  # give the server start up time
+    async with AsyncSubstrateInterface(
+        "ws://localhost:8080",
+        ss58_format=42,
+        chain_name="Bittensor",
+        retry_timeout=10.0,
+        ws_shutdown_timer=None,
+    ) as substrate:
+        blocks_to_check = [
+            5215000,
+            5215001,
+            5215002,
+            5215003,
+            5215004,
+            5215005,
+            5215006,
+        ]
+        tasks = []
+        for block in blocks_to_check:
+            block_hash = await substrate.get_block_hash(block_id=block)
+            tasks.append(
+                substrate.query_map(
+                    "SubtensorModule", "TotalHotkeyShares", block_hash=block_hash
+                )
+            )
+        records = await asyncio.gather(*tasks)
+        assert len(records) == len(blocks_to_check)
+        await substrate.close()
+        with open(ws_logger_path, "r") as f:
+            assert "Pausing" in f.read()
+        with open(asi_logger_path, "r") as f:
+            assert "Timeout/ConnectionClosed occurred." in f.read()
+    shutdown_thread = threading.Thread(target=proxy.close)
+    shutdown_thread.start()
+    shutdown_thread.join(timeout=5)
+    server_thread.join(timeout=5)