99import logging
1010import ssl
1111import time
12+ import warnings
13+ from unittest .mock import AsyncMock
1214from hashlib import blake2b
1315from typing import (
1416 Optional ,
@@ -530,15 +532,31 @@ def __init__(
530532 self ._exit_task = None
531533 self ._open_subscriptions = 0
532534 self ._options = options if options else {}
533- self .last_received = time .time ()
535+ try :
536+ now = asyncio .get_running_loop ().time ()
537+ except RuntimeError :
538+ warnings .warn (
539+ "You are instantiating the AsyncSubstrateInterface Websocket outside of an event loop. "
540+ "Verify this is intended."
541+ )
542+ now = asyncio .new_event_loop ().time ()
543+ self .last_received = now
544+ self .last_sent = now
534545
535546 async def __aenter__ (self ):
536547 async with self ._lock :
537548 self ._in_use += 1
538549 await self .connect ()
539550 return self
540551
552+ @staticmethod
553+ async def loop_time () -> float :
554+ return asyncio .get_running_loop ().time ()
555+
541556 async def connect (self , force = False ):
557+ now = await self .loop_time ()
558+ self .last_received = now
559+ self .last_sent = now
542560 if self ._exit_task :
543561 self ._exit_task .cancel ()
544562 if not self ._initialized or force :
@@ -594,7 +612,7 @@ async def _recv(self) -> None:
594612 try :
595613 # TODO consider wrapping this in asyncio.wait_for and use that for the timeout logic
596614 response = json .loads (await self .ws .recv (decode = False ))
597- self .last_received = time . time ()
615+ self .last_received = await self . loop_time ()
598616 async with self ._lock :
599617 # note that these 'subscriptions' are all waiting sent messages which have not received
600618 # responses, and thus are not the same as RPC 'subscriptions', which are unique
@@ -630,12 +648,12 @@ async def send(self, payload: dict) -> int:
630648 Returns:
631649 id: the internal ID of the request (incremented int)
632650 """
633- # async with self._lock:
634651 original_id = get_next_id ()
635652 # self._open_subscriptions += 1
636653 await self .max_subscriptions .acquire ()
637654 try :
638655 await self .ws .send (json .dumps ({** payload , ** {"id" : original_id }}))
656+ self .last_sent = await self .loop_time ()
639657 return original_id
640658 except (ConnectionClosed , ssl .SSLError , EOFError ):
641659 async with self ._lock :
@@ -697,13 +715,16 @@ def __init__(
697715 self .chain_endpoint = url
698716 self .url = url
699717 self ._chain = chain_name
700- self .ws = Websocket (
701- url ,
702- options = {
703- "max_size" : self .ws_max_size ,
704- "write_limit" : 2 ** 16 ,
705- },
706- )
718+ if not _mock :
719+ self .ws = Websocket (
720+ url ,
721+ options = {
722+ "max_size" : self .ws_max_size ,
723+ "write_limit" : 2 ** 16 ,
724+ },
725+ )
726+ else :
727+ self .ws = AsyncMock (spec = Websocket )
707728 self ._lock = asyncio .Lock ()
708729 self .config = {
709730 "use_remote_preset" : use_remote_preset ,
@@ -726,9 +747,11 @@ def __init__(
726747 self ._initializing = False
727748 self .registry_type_map = {}
728749 self .type_id_to_name = {}
750+ self ._mock = _mock
729751
730752 async def __aenter__ (self ):
731- await self .initialize ()
753+ if not self ._mock :
754+ await self .initialize ()
732755 return self
733756
734757 async def initialize (self ):
@@ -2120,7 +2143,11 @@ async def _make_rpc_request(
21202143
21212144 if request_manager .is_complete :
21222145 break
2123- if time .time () - self .ws .last_received >= self .retry_timeout :
2146+ if (
2147+ (current_time := await self .ws .loop_time ()) - self .ws .last_received
2148+ >= self .retry_timeout
2149+ and current_time - self .ws .last_sent >= self .retry_timeout
2150+ ):
21242151 if attempt >= self .max_retries :
21252152 logger .warning (
21262153 f"Timed out waiting for RPC requests { attempt } times. Exiting."
0 commit comments