Suppress ConnectionLostError in shutdown to avoid dirty tracebacks that sometime include ExceptionGroups (#29)

Author: vrslev

pyproject.toml

@@ -66,7 +66,16 @@ ignore = [
     "S101",
     "SLF001",
     "CPY001",
+    "S104",
 ]
+extend-per-file-ignores = { "tests/*" = [
+    "ARG002",
+    "ARG003",
+    "PLR6301",
+    "ANN401",
+    "PLC2801",
+] }
+
 
 [tool.pytest.ini_options]
 addopts = "--cov -s"

stompman/client.py

@@ -1,6 +1,6 @@
 import asyncio
 from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable
-from contextlib import AsyncExitStack, asynccontextmanager
+from contextlib import AsyncExitStack, asynccontextmanager, suppress
 from dataclasses import dataclass, field
 from types import TracebackType
 from typing import Any, ClassVar, NamedTuple, Self, TypedDict
@@ -10,6 +10,7 @@
 from stompman.connection import AbstractConnection, Connection
 from stompman.errors import (
     ConnectionConfirmationTimeoutError,
+    ConnectionLostError,
     FailedAllConnectAttemptsError,
     UnsupportedProtocolVersionError,
 )
@@ -199,7 +200,10 @@ async def _connection_lifespan(self) -> AsyncGenerator[None, None]:
 
         async def send_heartbeats_forever() -> None:
             while True:
-                self._connection.write_heartbeat()
+                try:
+                    self._connection.write_heartbeat()
+                except ConnectionLostError:
+                    return
                 await asyncio.sleep(heartbeat_interval)
 
         async with asyncio.TaskGroup() as task_group:
@@ -209,10 +213,11 @@ async def send_heartbeats_forever() -> None:
             finally:
                 task.cancel()
 
-        await self._connection.write_frame(DisconnectFrame(headers={"receipt": str(uuid4())}))
-        await self._connection.read_frame_of_type(
-            ReceiptFrame, max_chunk_size=self.read_max_chunk_size, timeout=self.read_timeout
-        )
+        with suppress(ConnectionLostError):
+            await self._connection.write_frame(DisconnectFrame(headers={"receipt": str(uuid4())}))
+            await self._connection.read_frame_of_type(
+                ReceiptFrame, max_chunk_size=self.read_max_chunk_size, timeout=self.read_timeout
+            )
 
     @asynccontextmanager
     async def enter_transaction(self) -> AsyncGenerator[str, None]:

stompman/connection.py

@@ -1,7 +1,7 @@
 import asyncio
 import socket
 from collections.abc import AsyncGenerator, Generator, Iterator
-from contextlib import contextmanager
+from contextlib import contextmanager, suppress
 from dataclasses import dataclass
 from typing import Protocol, Self, TypeVar, cast
 
@@ -52,7 +52,7 @@ async def connect(cls, host: str, port: int, timeout: int) -> Self | None:
 
     async def close(self) -> None:
         self.writer.close()
-        with _reraise_connection_lost(ConnectionError):
+        with suppress(ConnectionError):
             await self.writer.wait_closed()
 
     def write_heartbeat(self) -> None:

testing/consumer.py

@@ -5,7 +5,7 @@
 
 async def main() -> None:
     async with (
-        stompman.Client(servers=[stompman.ConnectionParameters("0.0.0.0", 61616, "admin", "admin")]) as client,  # noqa: S104
+        stompman.Client(servers=[stompman.ConnectionParameters("0.0.0.0", 61616, "admin", "=123")]) as client,
         client.subscribe("DLQ"),
     ):
         async for event in client.listen():

testing/producer.py

@@ -5,7 +5,7 @@
 
 async def main() -> None:
     async with (
-        stompman.Client(servers=[stompman.ConnectionParameters("0.0.0.0", 61616, "admin", "admin")]) as client,  # noqa: S104
+        stompman.Client(servers=[stompman.ConnectionParameters("0.0.0.0", 61616, "admin", "admin")]) as client,
         client.enter_transaction() as transaction,
     ):
         for _ in range(10):

tests/integration.py

@@ -58,28 +58,19 @@ async def closed_client(server: stompman.ConnectionParameters) -> AsyncGenerator
         yield client
 
 
-async def test_raises_connection_lost_error_in_aexit(server: stompman.ConnectionParameters) -> None:
-    with pytest.raises(ConnectionLostError):
-        async with closed_client(server):
-            pass
+async def test_not_raises_connection_lost_error_in_aexit(server: stompman.ConnectionParameters) -> None:
+    async with closed_client(server):
+        pass
 
 
-async def test_raises_connection_lost_error_in_write_frame(server: stompman.ConnectionParameters) -> None:
-    client = await closed_client(server).__aenter__()  # noqa: PLC2801
-
-    with pytest.raises(ConnectionLostError):
-        await client._connection.write_frame(stompman.ConnectFrame(headers={"accept-version": "", "host": ""}))
-
-    with pytest.raises(ConnectionLostError):
-        await client.__aexit__(None, None, None)
+async def test_not_raises_connection_lost_error_in_write_frame(server: stompman.ConnectionParameters) -> None:
+    async with closed_client(server) as client:
+        with pytest.raises(ConnectionLostError):
+            await client._connection.write_frame(stompman.ConnectFrame(headers={"accept-version": "", "host": ""}))
 
 
 @pytest.mark.parametrize("anyio_backend", [("asyncio", {"use_uvloop": True})])
-async def test_raises_connection_lost_error_in_write_heartbeat(server: stompman.ConnectionParameters) -> None:
-    client = await closed_client(server).__aenter__()  # noqa: PLC2801
-
-    with pytest.raises(ConnectionLostError):
-        client._connection.write_heartbeat()
-
-    with pytest.raises(ConnectionLostError):
-        await client.__aexit__(None, None, None)
+async def test_not_raises_connection_lost_error_in_write_heartbeat(server: stompman.ConnectionParameters) -> None:
+    async with closed_client(server) as client:
+        with pytest.raises(ConnectionLostError):
+            client._connection.write_heartbeat()

tests/test_client.py

@@ -33,6 +33,7 @@
     UnsupportedProtocolVersionError,
 )
 from stompman.client import ConnectionParameters, ErrorEvent, HeartbeatEvent, MessageEvent
+from stompman.errors import ConnectionLostError
 
 pytestmark = pytest.mark.anyio
 
@@ -42,19 +43,19 @@ class BaseMockConnection(AbstractConnection):
     @classmethod
     async def connect(
         cls,
-        host: str,  # noqa: ARG003
-        port: int,  # noqa: ARG003
-        timeout: int,  # noqa: ARG003
+        host: str,
+        port: int,
+        timeout: int,
     ) -> Self | None:
         return cls()
 
     async def close(self) -> None: ...
     def write_heartbeat(self) -> None: ...
     async def write_frame(self, frame: AnyClientFrame) -> None: ...
-    async def read_frames(  # noqa: PLR6301
+    async def read_frames(
         self,
-        max_chunk_size: int,  # noqa: ARG002
-        timeout: int,  # noqa: ARG002
+        max_chunk_size: int,
+        timeout: int,
     ) -> AsyncGenerator[AnyServerFrame, None]:  # pragma: no cover
         await asyncio.Future()
         yield  # type: ignore[misc]
@@ -65,13 +66,13 @@ def create_spying_connection(
 ) -> tuple[type[AbstractConnection], list[AnyClientFrame | AnyServerFrame | HeartbeatFrame]]:
     @dataclass
     class BaseCollectingConnection(BaseMockConnection):
-        async def write_frame(self, frame: AnyClientFrame) -> None:  # noqa: PLR6301
+        async def write_frame(self, frame: AnyClientFrame) -> None:
             collected_frames.append(frame)
 
-        async def read_frames(  # noqa: PLR6301
+        async def read_frames(
             self,
-            max_chunk_size: int,  # noqa: ARG002
-            timeout: int,  # noqa: ARG002
+            max_chunk_size: int,
+            timeout: int,
         ) -> AsyncGenerator[AnyServerFrame, None]:
             for frame in next(read_frames_iterator):
                 collected_frames.append(frame)
@@ -147,9 +148,9 @@ class MockConnection(BaseMockConnection):
         @classmethod
         async def connect(
             cls,
-            host: str,  # noqa: ARG003
-            port: int,  # noqa: ARG003
-            timeout: int,  # noqa: ARG003
+            host: str,
+            port: int,
+            timeout: int,
         ) -> Self | None:
             return None
 
@@ -185,9 +186,9 @@ class MockConnection(BaseMockConnection):
         @classmethod
         async def connect(
             cls,
-            host: str,  # noqa: ARG003
-            port: int,  # noqa: ARG003
-            timeout: int,  # noqa: ARG003
+            host: str,
+            port: int,
+            timeout: int,
         ) -> Self | None:
             return None
 
@@ -241,7 +242,7 @@ 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) -> Any:  # noqa: ANN401
+    async def timeout(future: Awaitable[Any], timeout: float) -> Any:
         assert timeout == client.connection_confirmation_timeout
         return await original_wait_for(future, 0)
 
@@ -250,7 +251,7 @@ async def timeout(future: Awaitable[Any], timeout: float) -> Any:  # noqa: ANN40
 
     client = EnrichedClient(connection_class=BaseMockConnection)
     with pytest.raises(ConnectionConfirmationTimeoutError) as exc_info:
-        await client.__aenter__()  # noqa: PLC2801
+        await client.__aenter__()
 
     assert exc_info.value == ConnectionConfirmationTimeoutError(client.connection_confirmation_timeout)
 
@@ -263,7 +264,7 @@ async def test_client_lifespan_unsupported_protocol_version() -> None:
 
     client = EnrichedClient(connection_class=connection_class)
     with pytest.raises(UnsupportedProtocolVersionError) as exc_info:
-        await client.__aenter__()  # noqa: PLC2801
+        await client.__aenter__()
 
     assert exc_info.value == UnsupportedProtocolVersionError(
         given_version=given_version, supported_version=client.PROTOCOL_VERSION
@@ -319,6 +320,16 @@ class MockConnection(connection_class):  # type: ignore[valid-type, misc]
     assert write_heartbeat_mock.mock_calls == [mock.call(), mock.call(), mock.call()]
 
 
+async def test_client_heartbeat_not_raises_connection_lost() -> None:
+    connection_class, _ = create_spying_connection(get_read_frames_with_lifespan([]))
+
+    class MockConnection(connection_class):  # type: ignore[valid-type, misc]
+        write_heartbeat = mock.Mock(side_effect=ConnectionLostError)
+
+    async with EnrichedClient(connection_class=MockConnection):
+        await asyncio.sleep(0)
+
+
 async def test_client_listen_to_events_ok() -> None:
     message_frame = MessageFrame(headers={"destination": "", "message-id": "", "subscription": ""}, body=b"hello")
     error_frame = ErrorFrame(headers={"message": "short description"})
@@ -423,11 +434,11 @@ async def test_message_event_with_auto_ack_ack_raises() -> None:
     event, ack, nack, on_suppressed_exception = get_mocked_message_event()
 
     async def func() -> None:  # noqa: RUF029
-        raise Exception  # noqa: TRY002
+        raise ImportError
 
-    with suppress(Exception):
+    with suppress(ImportError):
         await event.with_auto_ack(
-            func(), supressed_exception_classes=(RuntimeError,), on_suppressed_exception=on_suppressed_exception
+            func(), supressed_exception_classes=(ModuleNotFoundError,), on_suppressed_exception=on_suppressed_exception
         )
 
     ack.assert_called_once_with()

tests/test_connection.py

@@ -19,8 +19,8 @@ async def make_connection() -> Connection | None:
 
 async def make_mocked_connection(
     monkeypatch: pytest.MonkeyPatch,
-    reader: Any,  # noqa: ANN401
-    writer: Any,  # noqa: ANN401
+    reader: Any,
+    writer: Any,
 ) -> Connection:
     monkeypatch.setattr("asyncio.open_connection", mock.AsyncMock(return_value=(reader, writer)))
     connection = await make_connection()
@@ -29,7 +29,7 @@ async def make_mocked_connection(
 
 
 def mock_wait_for(monkeypatch: pytest.MonkeyPatch) -> None:
-    async def mock_impl(future: Awaitable[Any], timeout: int) -> Any:  # noqa: ANN401, ARG001
+    async def mock_impl(future: Awaitable[Any], timeout: int) -> Any:  # noqa: ARG001
         return await original_wait_for(future, timeout=0)
 
     original_wait_for = asyncio.wait_for
@@ -104,8 +104,7 @@ class MockWriter:
         wait_closed = mock.AsyncMock(side_effect=ConnectionError)
 
     connection = await make_mocked_connection(monkeypatch, mock.Mock(), MockWriter())
-    with pytest.raises(ConnectionLostError):
-        await connection.close()
+    await connection.close()
 
 
 async def test_connection_write_heartbeat_runtime_error(monkeypatch: pytest.MonkeyPatch) -> None: