Merge branch 'staging' into fix/thewhaleking/better-testing-maybe

Author: thewhaleking

async_substrate_interface/async_substrate.py

@@ -694,9 +694,17 @@ async def _cancel(self):
 
     async def connect(self, force=False):
         if not force:
-            await self._lock.acquire()
+            async with self._lock:
+                return await self._connect_internal(force)
         else:
             logger.debug("Proceeding without acquiring lock.")
+            return await self._connect_internal(force)
+
+    async def _connect_internal(self, force):
+        # Check state again after acquiring lock to avoid duplicate connections
+        if not force and self.state in (State.OPEN, State.CONNECTING):
+            return None
+
         logger.debug(f"Websocket connecting to {self.ws_url}")
         if self._sending is None or self._sending.empty():
             self._sending = asyncio.Queue()
@@ -725,17 +733,13 @@ async def connect(self, force=False):
             except socket.gaierror:
                 logger.debug(f"Hostname not known (this is just for testing")
                 await asyncio.sleep(10)
-                if self._lock.locked():
-                    self._lock.release()
                 return await self.connect(force=force)
             logger.debug("Connection established")
             self.ws = connection
             if self._send_recv_task is None or self._send_recv_task.done():
                 self._send_recv_task = asyncio.get_running_loop().create_task(
                     self._handler(self.ws)
                 )
-        if self._lock.locked():
-            self._lock.release()
         return None
 
     async def _handler(self, ws: ClientConnection) -> Union[None, Exception]:

tests/integration_tests/test_async_substrate_interface.py

@@ -302,3 +302,29 @@ async def test_get_payment_info():
         assert partial_fee_all_options > partial_fee_no_era
         assert partial_fee_all_options > partial_fee_era
     print("test_get_payment_info succeeded")
+
+
+@pytest.mark.asyncio
+async def test_concurrent_rpc_requests():
+    """
+    Test that multiple concurrent RPC requests on a shared connection work correctly.
+
+    This test verifies the fix for the issue where multiple concurrent tasks
+    re-initializing the WebSocket connection caused requests to hang.
+    """
+    print("Testing test_concurrent_rpc_requests")
+
+    async def concurrent_task(substrate, task_id):
+        """Make multiple RPC calls from a single task."""
+        for i in range(5):
+            result = await substrate.get_block_number(None)
+            assert isinstance(result, int)
+            assert result > 0
+
+    async with AsyncSubstrateInterface(LATENT_LITE_ENTRYPOINT) as substrate:
+        # Run 5 concurrent tasks, each making 5 RPC calls (25 total)
+        # This tests that the connection is properly shared without re-initialization
+        tasks = [concurrent_task(substrate, i) for i in range(5)]
+        await asyncio.gather(*tasks)
+
+    print("test_concurrent_rpc_requests succeeded")