Improve connection loss handling (#19)

Author: vrslev

stompman/connection.py

@@ -1,6 +1,7 @@
 import asyncio
 import socket
-from collections.abc import AsyncGenerator, Iterator
+from collections.abc import AsyncGenerator, Generator, Iterator
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from typing import Protocol, Self, TypedDict, TypeVar, cast
 
@@ -9,7 +10,7 @@
 from stompman.protocol import NEWLINE, Parser, dump_frame
 
 
-class _MultiHostHostLike(TypedDict):
+class MultiHostHostLike(TypedDict):
     username: str | None
     password: str | None
     host: str | None
@@ -24,7 +25,7 @@ class ConnectionParameters:
     passcode: str = field(repr=False)
 
     @classmethod
-    def from_pydantic_multihost_hosts(cls, hosts: list[_MultiHostHostLike]) -> list[Self]:
+    def from_pydantic_multihost_hosts(cls, hosts: list[MultiHostHostLike]) -> list[Self]:
         """Create connection parameters from a list of `MultiHostUrl` objects.
 
         .. code-block:: python
@@ -101,16 +102,25 @@ async def connect(self) -> bool:
             return False
         return True
 
+    @contextmanager
+    def _reraise_connection_lost(self, *causes: type[Exception]) -> Generator[None, None, None]:
+        try:
+            yield
+        except causes as exception:
+            raise ConnectionLostError(self.read_timeout) from exception
+
     async def close(self) -> None:
         self.writer.close()
-        await self.writer.wait_closed()
+        with self._reraise_connection_lost(ConnectionError):
+            await self.writer.wait_closed()
 
     def write_heartbeat(self) -> None:
         return self.writer.write(NEWLINE)
 
     async def write_frame(self, frame: AnyClientFrame) -> None:
         self.writer.write(dump_frame(frame))
-        await self.writer.drain()
+        with self._reraise_connection_lost(ConnectionError):
+            await self.writer.drain()
 
     async def _read_non_empty_bytes(self) -> bytes:
         while (
@@ -123,10 +133,8 @@ async def read_frames(self) -> AsyncGenerator[AnyServerFrame, None]:
         parser = Parser()
 
         while True:
-            try:
+            with self._reraise_connection_lost(ConnectionError, TimeoutError):
                 raw_frames = await asyncio.wait_for(self._read_non_empty_bytes(), timeout=self.read_timeout)
-            except (TimeoutError, ConnectionError) as exception:
-                raise ConnectionLostError(self.read_timeout) from exception
 
             for frame in cast(Iterator[AnyServerFrame], parser.load_frames(raw_frames)):
                 yield frame

stompman/frames.py

@@ -123,7 +123,7 @@ class NackFrame:
     body: bytes = b""
 
 
-BeginHeaders = TypedDict("BeginHeaders", {"transaction": NotRequired[str], "content-length": NotRequired[str]})
+BeginHeaders = TypedDict("BeginHeaders", {"transaction": str, "content-length": NotRequired[str]})
 
 
 @dataclass
@@ -132,7 +132,7 @@ class BeginFrame:
     body: bytes = b""
 
 
-CommitHeaders = TypedDict("CommitHeaders", {"transaction": NotRequired[str], "content-length": NotRequired[str]})
+CommitHeaders = TypedDict("CommitHeaders", {"transaction": str, "content-length": NotRequired[str]})
 
 
 @dataclass
@@ -141,7 +141,7 @@ class CommitFrame:
     body: bytes = b""
 
 
-AbortHeaders = TypedDict("AbortHeaders", {"transaction": NotRequired[str], "content-length": NotRequired[str]})
+AbortHeaders = TypedDict("AbortHeaders", {"transaction": str, "content-length": NotRequired[str]})
 
 
 @dataclass

tests/test_connection.py

@@ -1,4 +1,5 @@
 import asyncio
+import socket
 from collections.abc import Awaitable
 from functools import partial
 from typing import Any
@@ -14,7 +15,7 @@
     ConnectionParameters,
     HeartbeatFrame,
 )
-from stompman.frames import CommitFrame
+from stompman.frames import BeginFrame, CommitFrame
 
 
 @pytest.fixture()
@@ -110,13 +111,40 @@ async def take_frames(count: int) -> list[AnyServerFrame]:
     assert MockWriter.write.mock_calls == [mock.call(b"\n"), mock.call(b"COMMIT\ntransaction:transaction\n\n\x00")]
 
 
+async def test_connection_close_connection_error(connection: Connection, monkeypatch: pytest.MonkeyPatch) -> None:
+    class MockWriter:
+        close = mock.Mock()
+        wait_closed = mock.AsyncMock(side_effect=ConnectionError)
+
+    monkeypatch.setattr("asyncio.open_connection", mock.AsyncMock(return_value=(mock.Mock(), MockWriter())))
+    await connection.connect()
+
+    with pytest.raises(ConnectionLostError):
+        await connection.close()
+
+
+async def test_connection_write_frame_connection_error(connection: Connection, monkeypatch: pytest.MonkeyPatch) -> None:
+    class MockWriter:
+        write = mock.Mock()
+        drain = mock.AsyncMock(side_effect=ConnectionError)
+
+    monkeypatch.setattr("asyncio.open_connection", mock.AsyncMock(return_value=(mock.Mock(), MockWriter())))
+    await connection.connect()
+
+    with pytest.raises(ConnectionLostError):
+        await connection.write_frame(BeginFrame(headers={"transaction": ""}))
+
+
 async def test_connection_timeout(monkeypatch: pytest.MonkeyPatch, connection: Connection) -> None:
     mock_wait_for(monkeypatch)
     assert not await connection.connect()
 
 
-async def test_connection_error(monkeypatch: pytest.MonkeyPatch, connection: Connection) -> None:
-    monkeypatch.setattr("asyncio.open_connection", mock.AsyncMock(side_effect=ConnectionError))
+@pytest.mark.parametrize("exception", [BrokenPipeError, socket.gaierror])
+async def test_connection_connect_connection_error(
+    monkeypatch: pytest.MonkeyPatch, connection: Connection, exception: type[Exception]
+) -> None:
+    monkeypatch.setattr("asyncio.open_connection", mock.AsyncMock(side_effect=exception))
     assert not await connection.connect()