feat: add timeout to provider builder

Author: Matthias Zimmermann

src/arkiv/client.py

@@ -247,14 +247,30 @@ async def is_connected(self, show_traceback: bool = False) -> bool:
 
     async def __aenter__(self) -> AsyncArkiv:
         """Enter async context manager."""
-        # Initialize pending account if provided
-        if self._pending_account:
-            await self._initialize_account_async(self._pending_account)
-            self._pending_account = None
+        try:
+            # Initialize pending account if provided
+            if self._pending_account:
+                await self._initialize_account_async(self._pending_account)
+                self._pending_account = None
+
+            # Populate connection cache
+            await self.is_connected()
+            return self
+        except Exception:
+            # Best-effort cleanup if entering the context fails
+            logger.debug(
+                "AsyncArkiv.__aenter__ failed, attempting cleanup before re-raising"
+            )
+            try:
+                await self.arkiv.cleanup_filters()
+            except Exception:
+                logger.exception(
+                    "Error while cleaning up filters after __aenter__ failure"
+                )
 
-        # Populate connection cache
-        await self.is_connected()
-        return self
+            await self._disconnect_provider()
+            self._cached_connected = False
+            raise
 
     async def __aexit__(
         self,
@@ -267,9 +283,8 @@ async def __aexit__(
         logger.debug("Cleaning up event filters...")
         await self.arkiv.cleanup_filters()
 
-        # Disconnect provider if it has disconnect method
-        if hasattr(self.provider, "disconnect"):
-            await self.provider.disconnect()
+        # Disconnect provider and close underlying resources
+        await self._disconnect_provider()
 
         # Then stop the node if managed and update cache
         self._cleanup_node()
@@ -299,6 +314,18 @@ async def _initialize_account_async(self, account: NamedAccount) -> None:
             f"Account balance for {account.name} ({account.address}): {balance_eth} ETH"
         )
 
+    async def _disconnect_provider(self) -> None:
+        """Best-effort async disconnect of the underlying provider, if supported."""
+        provider = self.provider
+        if provider is None:
+            return
+
+        if hasattr(provider, "disconnect"):
+            try:
+                await provider.disconnect()  # type: ignore[func-returns-value]
+            except Exception:
+                logger.exception("Error while disconnecting async provider")
+
     def _cleanup_node(self) -> None:
         """Cleanup node and update connection cache."""
         super()._cleanup_node()

src/arkiv/provider.py

@@ -73,6 +73,7 @@ def __init__(self) -> None:
         self._transport: TransportType = cast(TransportType, TRANSPORT_DEFAULT)
         self._port: int | None = DEFAULT_PORT  # Set default port for localhost
         self._url: str | None = None
+        self._timeout_in: int | None = None  # timeout in seconds
         self._node: ArkivNode | None = None
         self._is_async: bool = False  # Default to sync providers
 
@@ -103,7 +104,7 @@ def kaolin(self) -> ProviderBuilder:
         self._port = None
         return self
 
-    def custom(self, url: str) -> ProviderBuilder:
+    def custom(self, url: str, fallback_url: str | None = None) -> ProviderBuilder:
         """
         Configure with custom RPC URL.
 
@@ -186,6 +187,16 @@ def ws(self) -> ProviderBuilder:
         self._transport = cast(TransportType, WS)
         return self
 
+    def timeout(self, seconds: int) -> ProviderBuilder:
+        """
+        Sets the request timeout for the provider.
+
+        Args:
+            seconds: Timeout duration in seconds
+        """
+        self._timeout_in = seconds
+        return self
+
     def async_mode(self, async_provider: bool = True) -> ProviderBuilder:
         """
         Sets the async provider mode.
@@ -276,9 +287,32 @@ def build(self) -> BaseProvider | AsyncBaseProvider:
         if self._transport == HTTP:
             # Consider async mode
             if self._is_async:
-                return AsyncHTTPProvider(url)
+                if self._timeout_in is not None:
+                    import aiohttp
+
+                    timeout = aiohttp.ClientTimeout(total=self._timeout_in)
+                    return AsyncHTTPProvider(url, request_kwargs={"timeout": timeout})
+                else:
+                    return AsyncHTTPProvider(url)
             else:
-                return HTTPProvider(url)
+                if self._timeout_in is not None:
+                    return HTTPProvider(
+                        url, request_kwargs={"timeout": self._timeout_in}
+                    )
+                else:
+                    return HTTPProvider(url)
         # Web socket transport (always async)
         else:
+            if self._timeout_in is not None:
+                return cast(
+                    AsyncBaseProvider,
+                    WebSocketProvider(
+                        url,
+                        request_timeout=self._timeout_in,
+                        # websocket_kwargs={
+                        #     "ping_interval": self._timeout_in,
+                        #     "ping_timeout": self._timeout_in * 2,
+                        # },
+                    ),
+                )
             return cast(AsyncBaseProvider, WebSocketProvider(url))

tests/conftest.py

@@ -4,8 +4,12 @@
 Provides either external node connections (via env vars) or containerized test nodes.
 """
 
+import json
 import logging
 import os
+import socketserver
+import threading
+import time
 from collections.abc import AsyncGenerator, Generator
 from pathlib import Path
 
@@ -27,6 +31,10 @@
 WALLET_FILE_ENV_PREFIX = "WALLET_FILE"
 WALLET_PASSWORD_ENV_PREFIX = "WALLET_PASSWORD"
 
+SLOW_LOCAL_SERVER_HOST = "127.0.0.1"
+SLOW_LOCAL_SERVER_PORT = 9876
+SLOW_LOCAL_SERVER_TIMEOUT = 5
+
 ALICE = "alice"
 BOB = "bob"
 
@@ -50,6 +58,40 @@ def _load_env_if_available() -> None:
 _load_env_if_available()
 
 
+@pytest.fixture(scope="session")
+def delayed_rpc_server():
+    class DelayedHandler(socketserver.BaseRequestHandler):
+        def handle(self):
+            time.sleep(SLOW_LOCAL_SERVER_TIMEOUT)  # Delay > timeout
+            self.request.sendall(
+                json.dumps({"jsonrpc": "2.0", "id": 1, "result": 42}).encode()
+            )
+
+    httpd = socketserver.TCPServer(
+        (SLOW_LOCAL_SERVER_HOST, SLOW_LOCAL_SERVER_PORT), DelayedHandler
+    )
+
+    httpd.allow_reuse_address = True
+
+    server_thread = threading.Thread(target=httpd.serve_forever, daemon=True)
+    server_thread.start()
+
+    # Wait & verify server is listening
+    time.sleep(1.0)
+    import socket
+
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    assert sock.connect_ex((SLOW_LOCAL_SERVER_HOST, SLOW_LOCAL_SERVER_PORT)) == 0, (
+        "Server not listening"
+    )
+    sock.close()
+
+    yield f"http://{SLOW_LOCAL_SERVER_HOST}:{SLOW_LOCAL_SERVER_PORT}"
+
+    httpd.shutdown()
+    httpd.server_close()
+
+
 @pytest.fixture(scope="session")
 def arkiv_container() -> Generator[ArkivNode, None, None]:
     """