Show received frames in ConnectionConfirmationTimeoutError (#38)

Author: vrslev

Justfile

@@ -5,7 +5,7 @@ install:
     uv -q sync
 
 test *args:
-    uv -q run pytest {{args}}
+    .venv/bin/pytest {{args}}
 
 lint:
     uv -q run ruff check .

stompman/client.py

@@ -160,6 +160,26 @@ async def _connect_to_any_server(self) -> None:
             timeout=self.connect_timeout,
         )
 
+    async def _wait_for_connected_frame(self) -> ConnectedFrame:
+        collected_frames = []
+
+        async def inner() -> ConnectedFrame:
+            async for frame in self._connection.read_frames(
+                max_chunk_size=self.read_max_chunk_size, timeout=self.read_timeout
+            ):
+                if isinstance(frame, ConnectedFrame):
+                    return frame
+                collected_frames.append(frame)
+            msg = "unreachable"  # pragma: no cover
+            raise AssertionError(msg)  # pragma: no cover
+
+        try:
+            return await asyncio.wait_for(inner(), timeout=self.connection_confirmation_timeout)
+        except TimeoutError as exception:
+            raise ConnectionConfirmationTimeoutError(
+                timeout=self.connection_confirmation_timeout, frames=collected_frames
+            ) from exception
+
     @asynccontextmanager
     async def _connection_lifespan(self) -> AsyncGenerator[None, None]:
         # On startup:
@@ -182,15 +202,7 @@ async def _connection_lifespan(self) -> AsyncGenerator[None, None]:
                 },
             )
         )
-        try:
-            connected_frame = await asyncio.wait_for(
-                self._connection.read_frame_of_type(
-                    ConnectedFrame, max_chunk_size=self.read_max_chunk_size, timeout=self.read_timeout
-                ),
-                timeout=self.connection_confirmation_timeout,
-            )
-        except TimeoutError as exception:
-            raise ConnectionConfirmationTimeoutError(self.connection_confirmation_timeout) from exception
+        connected_frame = await self._wait_for_connected_frame()
 
         if connected_frame.headers["version"] != self.PROTOCOL_VERSION:
             raise UnsupportedProtocolVersionError(
@@ -222,9 +234,11 @@ async def send_heartbeats_forever() -> None:
         if self._connection.active:
             await self._connection.write_frame(DisconnectFrame(headers={"receipt": str(uuid4())}))
         if self._connection.active:
-            await self._connection.read_frame_of_type(
-                ReceiptFrame, max_chunk_size=self.read_max_chunk_size, timeout=self.read_timeout
-            )
+            async for frame in self._connection.read_frames(
+                max_chunk_size=self.read_max_chunk_size, timeout=self.read_timeout
+            ):
+                if isinstance(frame, ReceiptFrame):
+                    break
 
     @asynccontextmanager
     async def enter_transaction(self) -> AsyncGenerator[str, None]:

stompman/connection.py

@@ -23,12 +23,6 @@ def write_heartbeat(self) -> None: ...
     async def write_frame(self, frame: AnyClientFrame) -> None: ...
     def read_frames(self, max_chunk_size: int, timeout: int) -> AsyncGenerator[AnyServerFrame, None]: ...
 
-    async def read_frame_of_type(self, type_: type[FrameType], max_chunk_size: int, timeout: int) -> FrameType:
-        while True:
-            async for frame in self.read_frames(max_chunk_size=max_chunk_size, timeout=timeout):
-                if isinstance(frame, type_):
-                    return frame
-
 
 @dataclass(kw_only=True)
 class Connection(AbstractConnection):

stompman/errors.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
+from stompman.frames import ErrorFrame, MessageFrame, ReceiptFrame
+
 if TYPE_CHECKING:
     from stompman.client import ConnectionParameters
 
@@ -14,6 +16,7 @@ def __str__(self) -> str:
 @dataclass
 class ConnectionConfirmationTimeoutError(Error):
     timeout: int
+    frames: list[MessageFrame | ReceiptFrame | ErrorFrame]
 
 
 @dataclass

tests/test_client.py

@@ -1,5 +1,5 @@
 import asyncio
-from collections.abc import AsyncGenerator, Awaitable
+from collections.abc import AsyncGenerator, Coroutine
 from contextlib import suppress
 from dataclasses import dataclass, field
 from typing import Any, Self
@@ -227,18 +227,28 @@ class MockConnection(connection_class):  # type: ignore[valid-type, misc]
 
 
 async def test_client_lifespan_connection_not_confirmed(monkeypatch: pytest.MonkeyPatch) -> None:
-    async def timeout(future: Awaitable[Any], timeout: float) -> object:
+    async def timeout(future: Coroutine[Any, Any, Any], timeout: float) -> object:
         assert timeout == client.connection_confirmation_timeout
-        return await original_wait_for(future, 0)
+        task = asyncio.create_task(future)
+        await asyncio.sleep(0)
+        return await original_wait_for(task, 0)
 
     original_wait_for = asyncio.wait_for
     monkeypatch.setattr("asyncio.wait_for", timeout)
 
-    client = EnrichedClient(connection_class=BaseMockConnection)
+    class MockConnection(BaseMockConnection):
+        @staticmethod
+        async def read_frames(max_chunk_size: int, timeout: int) -> AsyncGenerator[AnyServerFrame, None]:
+            yield ErrorFrame(headers={"message": "hi"})
+            await asyncio.sleep(0)
+
+    client = EnrichedClient(connection_class=MockConnection)
     with pytest.raises(ConnectionConfirmationTimeoutError) as exc_info:
         await client.__aenter__()  # noqa: PLC2801
 
-    assert exc_info.value == ConnectionConfirmationTimeoutError(client.connection_confirmation_timeout)
+    assert exc_info.value == ConnectionConfirmationTimeoutError(
+        timeout=client.connection_confirmation_timeout, frames=[ErrorFrame(headers={"message": "hi"})]
+    )
 
 
 async def test_client_lifespan_unsupported_protocol_version() -> None:

tests/test_errors.py

@@ -2,5 +2,5 @@
 
 
 def test_error_str() -> None:
-    error = ConnectionConfirmationTimeoutError(timeout=1)
+    error = ConnectionConfirmationTimeoutError(timeout=1, frames=[])
     assert str(error) == repr(error)