fix: support WebSocket timeouts on Python 3.10

src/api_ws_suite/ws_probe.py

@@ -22,13 +22,17 @@ async def probe(
     if expected_messages < 1:
         raise ValueError("expected_messages must be at least 1")
 
+    async def receive_messages():
+        async with websockets.connect(uri) as socket:
+            return [await socket.recv() for _ in range(expected_messages)]
+
     failures: list[str] = []
     for attempt in range(1, attempts + 1):
         started = time.perf_counter()
         try:
-            async with asyncio.timeout(timeout_seconds):
-                async with websockets.connect(uri) as socket:
-                    raw_messages = [await socket.recv() for _ in range(expected_messages)]
+            raw_messages = await asyncio.wait_for(
+                receive_messages(), timeout=timeout_seconds
+            )
             latency_ms = round((time.perf_counter() - started) * 1_000, 2)
             events: list[Any] = []
             for message_index, raw in enumerate(raw_messages):
@@ -58,7 +62,7 @@ async def probe(
                 "latency_ms": latency_ms,
                 "errors": errors,
             }
-        except (OSError, TimeoutError, websockets.WebSocketException) as exc:
+        except (OSError, asyncio.TimeoutError, websockets.WebSocketException) as exc:
             failures.append(f"attempt {attempt}: {exc.__class__.__name__}")
     return {
         "passed": False,

tests/test_ws_probe.py

@@ -2,7 +2,9 @@
 import json
 import socket
 
+import pytest
 from websockets.asyncio.server import serve
+from websockets.exceptions import ConnectionClosed
 
 from api_ws_suite.ws_probe import probe
 
@@ -63,6 +65,52 @@ def assertion(result):
     asyncio.run(run_server(handler, assertion, attempts=1, timeout_seconds=0.03))
 
 
+def test_timeout_budget_is_shared_across_messages():
+    async def handler(websocket):
+        await asyncio.sleep(0.08)
+        await websocket.send(json.dumps(event()))
+        await asyncio.sleep(0.08)
+        try:
+            await websocket.send(json.dumps(event("evt-2", 2)))
+        except ConnectionClosed:
+            # The probe may close the connection when its total budget expires.
+            pass
+
+    def assertion(result):
+        assert result["passed"] is False
+        assert result["errors"] == ["attempt 1: TimeoutError"]
+
+    asyncio.run(
+        run_server(
+            handler, assertion, attempts=1, timeout_seconds=0.12,
+            expected_messages=2,
+        )
+    )
+
+
+def test_cancelling_probe_propagates_without_retry():
+    async def run():
+        connected = asyncio.Event()
+        connections = 0
+
+        async def handler(websocket):
+            nonlocal connections
+            connections += 1
+            connected.set()
+            await websocket.wait_closed()
+
+        async with serve(handler, "127.0.0.1", 0) as server:
+            port = server.sockets[0].getsockname()[1]
+            task = asyncio.create_task(probe(f"ws://127.0.0.1:{port}", attempts=2))
+            await asyncio.wait_for(connected.wait(), timeout=1.0)
+            task.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await task
+            assert connections == 1
+
+    asyncio.run(run())
+
+
 def test_transient_failure_reconnects_within_bounded_attempts():
     connections = 0