Refactor tests (#57)

Author: vrslev

pyproject.toml

@@ -67,7 +67,7 @@ ignore = [
     "PLC2801",
     "PLR0913",
 ]
-extend-per-file-ignores = { "tests/*" = ["S101", "SLF001", "ARG"] }
+extend-per-file-ignores = { "tests/*" = ["S101", "SLF001", "ARG", "PLR6301"] }
 
 [tool.pytest.ini_options]
 addopts = "--cov -s -vv"

stompman/__init__.py

@@ -1,7 +1,10 @@
 from stompman.client import Client
 from stompman.config import ConnectionParameters, Heartbeat
 from stompman.errors import (
+    AllServersUnavailable,
+    AnyConnectionIssue,
     ConnectionConfirmationTimeout,
+    ConnectionLost,
     ConnectionLostError,
     Error,
     FailedAllConnectAttemptsError,
@@ -38,7 +41,9 @@
     "AbortFrame",
     "AckFrame",
     "AckMode",
+    "AllServersUnavailable",
     "AnyClientFrame",
+    "AnyConnectionIssue",
     "AnyRealServerFrame",
     "AnyServerFrame",
     "BeginFrame",
@@ -47,6 +52,7 @@
     "ConnectFrame",
     "ConnectedFrame",
     "ConnectionConfirmationTimeout",
+    "ConnectionLost",
     "ConnectionLostError",
     "ConnectionParameters",
     "DisconnectFrame",

stompman/connection_lifespan.py

@@ -1,5 +1,5 @@
 import asyncio
-from collections.abc import Callable
+from collections.abc import AsyncIterable, Callable
 from contextlib import suppress
 from dataclasses import dataclass
 from typing import Protocol
@@ -9,6 +9,7 @@
 from stompman.connection import AbstractConnection
 from stompman.errors import ConnectionConfirmationTimeout, StompProtocolConnectionIssue, UnsupportedProtocolVersion
 from stompman.frames import (
+    AnyServerFrame,
     ConnectedFrame,
     ConnectFrame,
     DisconnectFrame,
@@ -22,6 +23,54 @@
 from stompman.transaction import ActiveTransactions, commit_pending_transactions
 
 
+async def take_connected_frame(
+    *, frames_iter: AsyncIterable[AnyServerFrame], connection_confirmation_timeout: int
+) -> ConnectedFrame | ConnectionConfirmationTimeout:
+    collected_frames = []
+
+    async def take_connected_frame_and_collect_other_frames() -> ConnectedFrame:
+        async for frame in frames_iter:
+            if isinstance(frame, ConnectedFrame):
+                return frame
+            collected_frames.append(frame)
+        msg = "unreachable"
+        raise AssertionError(msg)
+
+    try:
+        return await asyncio.wait_for(
+            take_connected_frame_and_collect_other_frames(), timeout=connection_confirmation_timeout
+        )
+    except TimeoutError:
+        return ConnectionConfirmationTimeout(timeout=connection_confirmation_timeout, frames=collected_frames)
+
+
+def check_stomp_protocol_version(
+    *, connected_frame: ConnectedFrame, supported_version: str
+) -> UnsupportedProtocolVersion | None:
+    if connected_frame.headers["version"] == supported_version:
+        return None
+    return UnsupportedProtocolVersion(
+        given_version=connected_frame.headers["version"], supported_version=supported_version
+    )
+
+
+def calculate_heartbeat_interval(*, connected_frame: ConnectedFrame, client_heartbeat: Heartbeat) -> float:
+    server_heartbeat = Heartbeat.from_header(connected_frame.headers["heart-beat"])
+    return max(client_heartbeat.will_send_interval_ms, server_heartbeat.want_to_receive_interval_ms) / 1000
+
+
+async def wait_for_receipt_frame(
+    *, frames_iter: AsyncIterable[AnyServerFrame], disconnect_confirmation_timeout: int
+) -> None:
+    async def inner() -> None:
+        async for frame in frames_iter:
+            if isinstance(frame, ReceiptFrame):
+                break
+
+    with suppress(TimeoutError):
+        await asyncio.wait_for(inner(), timeout=disconnect_confirmation_timeout)
+
+
 class AbstractConnectionLifespan(Protocol):
     async def enter(self) -> StompProtocolConnectionIssue | None: ...
     async def exit(self) -> None: ...
@@ -51,61 +100,48 @@ async def _establish_connection(self) -> StompProtocolConnectionIssue | None:
                 },
             )
         )
-        collected_frames = []
-
-        async def take_connected_frame_and_collect_other_frames() -> ConnectedFrame:
-            async for frame in self.connection.read_frames():
-                if isinstance(frame, ConnectedFrame):
-                    return frame
-                collected_frames.append(frame)
-            msg = "unreachable"  # pragma: no cover
-            raise AssertionError(msg)  # pragma: no cover
-
-        try:
-            connected_frame = await asyncio.wait_for(
-                take_connected_frame_and_collect_other_frames(), timeout=self.connection_confirmation_timeout
-            )
-        except TimeoutError:
-            return ConnectionConfirmationTimeout(timeout=self.connection_confirmation_timeout, frames=collected_frames)
+        connected_frame_or_error = await take_connected_frame(
+            frames_iter=self.connection.read_frames(),
+            connection_confirmation_timeout=self.connection_confirmation_timeout,
+        )
+        if isinstance(connected_frame_or_error, ConnectionConfirmationTimeout):
+            return connected_frame_or_error
+        connected_frame = connected_frame_or_error
 
-        if connected_frame.headers["version"] != self.protocol_version:
-            return UnsupportedProtocolVersion(
-                given_version=connected_frame.headers["version"], supported_version=self.protocol_version
-            )
+        if unsupported_protocol_version_error := check_stomp_protocol_version(
+            connected_frame=connected_frame, supported_version=self.protocol_version
+        ):
+            return unsupported_protocol_version_error
 
-        server_heartbeat = Heartbeat.from_header(connected_frame.headers["heart-beat"])
         self.set_heartbeat_interval(
-            max(self.client_heartbeat.will_send_interval_ms, server_heartbeat.want_to_receive_interval_ms) / 1000
+            calculate_heartbeat_interval(connected_frame=connected_frame, client_heartbeat=self.client_heartbeat)
         )
         return None
 
     async def enter(self) -> StompProtocolConnectionIssue | None:
-        if connection_issue := await self._establish_connection():
-            return connection_issue
+        if protocol_connection_issue := await self._establish_connection():
+            return protocol_connection_issue
+
         await resubscribe_to_active_subscriptions(
             connection=self.connection, active_subscriptions=self.active_subscriptions
         )
         await commit_pending_transactions(connection=self.connection, active_transactions=self.active_transactions)
         return None
 
-    async def _take_receipt_frame(self) -> None:
-        async for frame in self.connection.read_frames():
-            if isinstance(frame, ReceiptFrame):
-                break
-
     async def exit(self) -> None:
         await unsubscribe_from_all_active_subscriptions(active_subscriptions=self.active_subscriptions)
         await self.connection.write_frame(DisconnectFrame(headers={"receipt": _make_receipt_id()}))
-
-        with suppress(TimeoutError):
-            await asyncio.wait_for(self._take_receipt_frame(), timeout=self.disconnect_confirmation_timeout)
-
-
-def _make_receipt_id() -> str:
-    return str(uuid4())
+        await wait_for_receipt_frame(
+            frames_iter=self.connection.read_frames(),
+            disconnect_confirmation_timeout=self.disconnect_confirmation_timeout,
+        )
 
 
 class ConnectionLifespanFactory(Protocol):
     def __call__(
         self, *, connection: AbstractConnection, connection_parameters: ConnectionParameters
     ) -> AbstractConnectionLifespan: ...
+
+
+def _make_receipt_id() -> str:
+    return str(uuid4())

stompman/connection_manager.py

@@ -1,5 +1,5 @@
 import asyncio
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Awaitable, Callable
 from dataclasses import dataclass, field
 from types import TracebackType
 from typing import TYPE_CHECKING, Self
@@ -26,6 +26,48 @@ class ActiveConnectionState:
     lifespan: "AbstractConnectionLifespan"
 
 
+async def attempt_to_connect(
+    *,
+    connect: Callable[[], Awaitable[ActiveConnectionState | AnyConnectionIssue]],
+    connect_retry_interval: int,
+    connect_retry_attempts: int,
+) -> ActiveConnectionState:
+    connection_issues = []
+
+    for attempt in range(connect_retry_attempts):
+        connection_result = await connect()
+        if isinstance(connection_result, ActiveConnectionState):
+            return connection_result
+
+        connection_issues.append(connection_result)
+        await asyncio.sleep(connect_retry_interval * (attempt + 1))
+
+    raise FailedAllConnectAttemptsError(retry_attempts=connect_retry_attempts, issues=connection_issues)
+
+
+async def connect_to_first_server(
+    connect_awaitables: list[Awaitable[ActiveConnectionState | None]],
+) -> ActiveConnectionState | None:
+    for maybe_connection_future in asyncio.as_completed(connect_awaitables):
+        if connection_state := await maybe_connection_future:
+            return connection_state
+    return None
+
+
+async def make_healthy_connection(
+    *, active_connection_state: ActiveConnectionState | None, servers: list[ConnectionParameters], connect_timeout: int
+) -> ActiveConnectionState | AnyConnectionIssue:
+    if not active_connection_state:
+        return AllServersUnavailable(servers=servers, timeout=connect_timeout)
+
+    try:
+        connection_issue = await active_connection_state.lifespan.enter()
+    except ConnectionLostError:
+        return ConnectionLost()
+
+    return active_connection_state if connection_issue is None else connection_issue
+
+
 @dataclass(kw_only=True, slots=True)
 class ConnectionManager:
     servers: list[ConnectionParameters]
@@ -70,47 +112,27 @@ async def _create_connection_to_one_server(self, server: ConnectionParameters) -
             )
         return None
 
-    async def _create_connection_to_any_server(self) -> ActiveConnectionState | None:
-        for maybe_connection_future in asyncio.as_completed(
-            [self._create_connection_to_one_server(server) for server in self.servers]
-        ):
-            if connection_state := await maybe_connection_future:
-                return connection_state
-        return None
-
     async def _connect_to_any_server(self) -> ActiveConnectionState | AnyConnectionIssue:
-        if not (active_connection_state := await self._create_connection_to_any_server()):
-            return AllServersUnavailable(servers=self.servers, timeout=self.connect_timeout)
-
-        try:
-            if connection_issue := await active_connection_state.lifespan.enter():
-                return connection_issue
-        except ConnectionLostError:
-            return ConnectionLost()
-
-        return active_connection_state
+        active_connection_state = await connect_to_first_server(
+            [self._create_connection_to_one_server(server) for server in self.servers]
+        )
+        return await make_healthy_connection(
+            active_connection_state=active_connection_state, servers=self.servers, connect_timeout=self.connect_timeout
+        )
 
     async def _get_active_connection_state(self) -> ActiveConnectionState:
         if self._active_connection_state:
             return self._active_connection_state
 
-        connection_issues: list[AnyConnectionIssue] = []
-
         async with self._reconnect_lock:
             if self._active_connection_state:
                 return self._active_connection_state
-
-            for attempt in range(self.connect_retry_attempts):
-                connection_result = await self._connect_to_any_server()
-
-                if isinstance(connection_result, ActiveConnectionState):
-                    self._active_connection_state = connection_result
-                    return connection_result
-
-                connection_issues.append(connection_result)
-                await asyncio.sleep(self.connect_retry_interval * (attempt + 1))
-
-        raise FailedAllConnectAttemptsError(retry_attempts=self.connect_retry_attempts, issues=connection_issues)
+            self._active_connection_state = await attempt_to_connect(
+                connect=self._connect_to_any_server,
+                connect_retry_interval=self.connect_retry_interval,
+                connect_retry_attempts=self.connect_retry_attempts,
+            )
+            return self._active_connection_state
 
     def _clear_active_connection_state(self) -> None:
         self._active_connection_state = None

tests/conftest.py

@@ -12,14 +12,9 @@
 from stompman.connection_manager import ConnectionManager
 
 
-@pytest.fixture(
-    params=[
-        pytest.param(("asyncio", {"use_uvloop": True}), id="asyncio+uvloop"),
-        pytest.param(("asyncio", {"use_uvloop": False}), id="asyncio"),
-    ],
-)
-def anyio_backend(request: pytest.FixtureRequest) -> object:
-    return request.param
+@pytest.fixture
+def anyio_backend() -> object:
+    return "asyncio"
 
 
 @pytest.fixture
@@ -127,11 +122,6 @@ async def read_frames() -> AsyncGenerator[stompman.AnyServerFrame, None]:
 CONNECTED_FRAME = stompman.ConnectedFrame(headers={"version": stompman.Client.PROTOCOL_VERSION, "heart-beat": "1,1"})
 
 
-@pytest.fixture(autouse=True)
-def _mock_receipt_id(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr(stompman.connection_lifespan, "_make_receipt_id", lambda: "receipt-id-1")
-
-
 def get_read_frames_with_lifespan(*read_frames: list[stompman.AnyServerFrame]) -> list[list[stompman.AnyServerFrame]]:
     return [
         [CONNECTED_FRAME],

tests/integration.py

@@ -21,13 +21,23 @@
     parse_header,
 )
 
-pytestmark = pytest.mark.anyio
-
 CONNECTION_PARAMETERS: Final = stompman.ConnectionParameters(
     host=os.environ["ARTEMIS_HOST"], port=61616, login="admin", passcode=":=123"
 )
 DESTINATION: Final = "DLQ"
 
+pytestmark = pytest.mark.anyio
+
+
+@pytest.fixture(
+    params=[
+        pytest.param(("asyncio", {"use_uvloop": True}), id="asyncio+uvloop"),
+        pytest.param(("asyncio", {"use_uvloop": False}), id="asyncio"),
+    ],
+)
+def anyio_backend(request: pytest.FixtureRequest) -> object:
+    return request.param
+
 
 @asynccontextmanager
 async def create_client() -> AsyncGenerator[stompman.Client, None]:

tests/test_client.py

@@ -0,0 +1,29 @@
+import asyncio
+from unittest import mock
+
+import pytest
+
+from tests.conftest import EnrichedClient, create_spying_connection, get_read_frames_with_lifespan
+
+pytestmark = pytest.mark.anyio
+
+
+async def test_client_heartbeats_ok(monkeypatch: pytest.MonkeyPatch) -> None:
+    async def mock_sleep(delay: float) -> None:
+        await real_sleep(0)
+        sleep_calls.append(delay)
+
+    sleep_calls: list[float] = []
+    real_sleep = asyncio.sleep
+    monkeypatch.setattr("asyncio.sleep", mock_sleep)
+
+    connection_class, _ = create_spying_connection(*get_read_frames_with_lifespan([]))
+    connection_class.write_heartbeat = (write_heartbeat_mock := mock.Mock())  # type: ignore[method-assign]
+
+    async with EnrichedClient(connection_class=connection_class):
+        await real_sleep(0)
+        await real_sleep(0)
+        await real_sleep(0)
+
+    assert sleep_calls == [0, 1, 1]
+    assert write_heartbeat_mock.mock_calls == [mock.call(), mock.call(), mock.call()]