Bug fixes

Author: thewhaleking

async_substrate_interface/async_substrate.py

@@ -586,6 +586,7 @@ def __init__(
         self._max_retries = max_retries
         self._last_activity = asyncio.Event()
         self._last_activity.set()
+        self._waiting_for_response = 0
 
     @property
     def state(self):
@@ -599,6 +600,14 @@ async def __aenter__(self):
             await self.connect()
         return self
 
+    async def mark_waiting_for_response(self):
+        async with self._lock:
+            self._waiting_for_response += 1
+
+    async def mark_response_received(self):
+        async with self._lock:
+            self._waiting_for_response -= 1
+
     @staticmethod
     async def loop_time() -> float:
         return asyncio.get_running_loop().time()
@@ -793,7 +802,12 @@ async def _handler(self, ws: ClientConnection) -> Union[None, Exception]:
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         if self.shutdown_timer is not None:
-            if self.state != State.CONNECTING:
+            if (
+                self.state != State.CONNECTING
+                and self._sending.qsize() == 0
+                and not self._received_subscriptions
+                and self._waiting_for_response <= 0
+            ):
                 if self._exit_task is not None:
                     self._exit_task.cancel()
                     try:
@@ -812,6 +826,7 @@ async def _exit_with_timer(self):
         try:
             if self.shutdown_timer is not None:
                 await asyncio.sleep(self.shutdown_timer)
+            logger.debug("Exiting with timer")
             await self.shutdown()
         except asyncio.CancelledError:
             pass
@@ -2495,6 +2510,7 @@ async def _make_rpc_request(
                 logger.debug(
                     f"Submitted payload ID {payload['id']} with websocket ID {item_id}: {output_payload}"
                 )
+            await ws.mark_waiting_for_response()
 
             while True:
                 for item_id in request_manager.unresponded():
@@ -2552,6 +2568,7 @@ async def _make_rpc_request(
                             )
 
                 if request_manager.is_complete:
+                    await ws.mark_response_received()
                     break
                 else:
                     await asyncio.sleep(0.01)
@@ -3948,6 +3965,7 @@ async def result_handler(message: dict, subscription_id) -> tuple[dict, bool]:
                 }
 
                 if "finalized" in message_result and wait_for_finalization:
+                    logger.debug("Extrinsic finalized. Unsubscribing.")
                     async with self.ws as ws:
                         await ws.unsubscribe(subscription_id)
                     return {
@@ -3956,14 +3974,17 @@ async def result_handler(message: dict, subscription_id) -> tuple[dict, bool]:
                         "finalized": True,
                     }, True
                 elif (
-                    "inblock" in message_result
+                    any(x in message_result for x in ["inblock", "inBlock"])
                     and wait_for_inclusion
                     and not wait_for_finalization
                 ):
+                    logger.debug("Extrinsic included. Unsubscribing.")
                     async with self.ws as ws:
                         await ws.unsubscribe(subscription_id)
                     return {
-                        "block_hash": message_result["inblock"],
+                        "block_hash": message_result.get(
+                            "inblock", message_result.get("inBlock")
+                        ),
                         "extrinsic_hash": "0x{}".format(extrinsic.extrinsic_hash.hex()),
                         "finalized": False,
                     }, True