Add disconnect confirmation timeout (#46)

Author: vrslev

stompman/client.py

@@ -1,6 +1,6 @@
 import asyncio
 from collections.abc import AsyncGenerator, AsyncIterator, Callable, Coroutine
-from contextlib import AsyncExitStack, asynccontextmanager
+from contextlib import AsyncExitStack, asynccontextmanager, suppress
 from dataclasses import dataclass, field
 from types import TracebackType
 from typing import ClassVar, Self
@@ -31,13 +31,14 @@
 
 
 @asynccontextmanager
-async def connection_lifespan(
+async def connection_lifespan(  # noqa: PLR0913
     *,
     connection: AbstractConnection,
     connection_parameters: ConnectionParameters,
     protocol_version: str,
     client_heartbeat: Heartbeat,
     connection_confirmation_timeout: int,
+    disconnect_confirmation_timeout: int,
 ) -> AsyncIterator[float]:
     await connection.write_frame(
         ConnectFrame(
@@ -52,7 +53,7 @@ async def connection_lifespan(
     )
     collected_frames = []
 
-    async def take_connected_frame() -> ConnectedFrame:
+    async def take_connected_frame_and_collect_other_frames() -> ConnectedFrame:
         async for frame in connection.read_frames():
             if isinstance(frame, ConnectedFrame):
                 return frame
@@ -61,7 +62,9 @@ async def take_connected_frame() -> ConnectedFrame:
         raise AssertionError(msg)  # pragma: no cover
 
     try:
-        connected_frame = await asyncio.wait_for(take_connected_frame(), timeout=connection_confirmation_timeout)
+        connected_frame = await asyncio.wait_for(
+            take_connected_frame_and_collect_other_frames(), timeout=connection_confirmation_timeout
+        )
     except TimeoutError as exception:
         raise ConnectionConfirmationTimeoutError(
             timeout=connection_confirmation_timeout, frames=collected_frames
@@ -79,9 +82,14 @@ async def take_connected_frame() -> ConnectedFrame:
     yield heartbeat_interval
 
     await connection.write_frame(DisconnectFrame(headers={"receipt": _make_receipt_id()}))
-    async for frame in connection.read_frames():
-        if isinstance(frame, ReceiptFrame):
-            break
+
+    async def take_receipt_frame() -> None:
+        async for frame in connection.read_frames():
+            if isinstance(frame, ReceiptFrame):
+                break
+
+    with suppress(TimeoutError):
+        await asyncio.wait_for(take_receipt_frame(), timeout=disconnect_confirmation_timeout)
 
 
 def _make_receipt_id() -> str:
@@ -211,9 +219,11 @@ class Client:
     connect_retry_attempts: int = 3
     connect_retry_interval: int = 1
     connect_timeout: int = 2
-    connection_confirmation_timeout: int = 2
     read_timeout: int = 2
     read_max_chunk_size: int = 1024 * 1024
+    connection_confirmation_timeout: int = 2
+    disconnect_confirmation_timeout: int = 2
+
     connection_class: type[AbstractConnection] = Connection
 
     _connection_manager: ConnectionManager = field(init=False)
@@ -263,6 +273,7 @@ async def _lifespan(
             protocol_version=self.PROTOCOL_VERSION,
             client_heartbeat=self.heartbeat,
             connection_confirmation_timeout=self.connection_confirmation_timeout,
+            disconnect_confirmation_timeout=self.disconnect_confirmation_timeout,
         ) as heartbeat_interval:
             self._restart_heartbeat_task(heartbeat_interval)
             async with subscriptions_lifespan(connection=connection, active_subscriptions=self._active_subscriptions):

tests/test_client.py

@@ -68,14 +68,6 @@ async def read_frames() -> AsyncGenerator[AnyServerFrame, None]:
     return BaseCollectingConnection, collected_frames
 
 
-def get_read_frames_with_lifespan(*read_frames: list[AnyServerFrame]) -> list[list[AnyServerFrame]]:
-    return [
-        [ConnectedFrame(headers={"version": Client.PROTOCOL_VERSION, "heart-beat": "1,1"})],
-        *read_frames,
-        [ReceiptFrame(headers={"receipt-id": "receipt-id-1"})],
-    ]
-
-
 CONNECT_FRAME = ConnectFrame(
     headers={
         "accept-version": Client.PROTOCOL_VERSION,
@@ -88,6 +80,14 @@ def get_read_frames_with_lifespan(*read_frames: list[AnyServerFrame]) -> list[li
 CONNECTED_FRAME = ConnectedFrame(headers={"version": Client.PROTOCOL_VERSION, "heart-beat": "1,1"})
 
 
+def get_read_frames_with_lifespan(*read_frames: list[AnyServerFrame]) -> list[list[AnyServerFrame]]:
+    return [
+        [CONNECTED_FRAME],
+        *read_frames,
+        [ReceiptFrame(headers={"receipt-id": "receipt-id-1"})],
+    ]
+
+
 def enrich_expected_frames(*expected_frames: AnyClientFrame | AnyServerFrame) -> list[AnyClientFrame | AnyServerFrame]:
     return [
         CONNECT_FRAME,
@@ -172,6 +172,30 @@ async def test_client_connection_lifespan_unsupported_protocol_version() -> None
     )
 
 
+async def test_client_connection_lifespan_disconnect_not_confirmed(monkeypatch: pytest.MonkeyPatch) -> None:
+    wait_for_calls = []
+
+    async def mock_wait_for(future: Coroutine[Any, Any, Any], timeout: float) -> object:
+        wait_for_calls.append(timeout)
+        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", mock_wait_for)
+    disconnect_confirmation_timeout = FAKER.pyint()
+    read_frames_yields = get_read_frames_with_lifespan([])
+    read_frames_yields[-1].clear()
+    connection_class, _ = create_spying_connection(*read_frames_yields)
+
+    async with EnrichedClient(
+        connection_class=connection_class, disconnect_confirmation_timeout=disconnect_confirmation_timeout
+    ):
+        pass
+
+    assert wait_for_calls[-1] == disconnect_confirmation_timeout
+
+
 async def test_client_heartbeats_ok(monkeypatch: pytest.MonkeyPatch) -> None:
     async def mock_sleep(delay: float) -> None:
         await real_sleep(0)