Refactor tests (second part) (#58)

Author: vrslev

stompman/client.py

@@ -76,7 +76,7 @@ def __post_init__(self) -> None:
     async def __aenter__(self) -> Self:
         self._task_group = await self._exit_stack.enter_async_context(asyncio.TaskGroup())
         self._heartbeat_task = self._task_group.create_task(asyncio.sleep(0))
-        await self._exit_stack.enter_async_context(self._connection_manager)
+        await self._connection_manager.enter()
         self._listen_task = self._task_group.create_task(self._listen_to_frames())
         return self
 
@@ -91,6 +91,7 @@ async def __aexit__(
             self._heartbeat_task.cancel()
             await asyncio.wait([self._listen_task, self._heartbeat_task])
             await self._exit_stack.aclose()
+            await self._connection_manager.exit()
 
     def _restart_heartbeat_task(self, interval: float) -> None:
         self._heartbeat_task.cancel()

stompman/connection_lifespan.py

@@ -1,8 +1,7 @@
 import asyncio
-from collections.abc import AsyncIterable, Callable
-from contextlib import suppress
-from dataclasses import dataclass
-from typing import Protocol
+from collections.abc import AsyncIterable, Awaitable, Callable
+from dataclasses import dataclass, field
+from typing import Any, Protocol, TypeVar
 from uuid import uuid4
 
 from stompman.config import ConnectionParameters, Heartbeat
@@ -22,26 +21,40 @@
 )
 from stompman.transaction import ActiveTransactions, commit_pending_transactions
 
+FrameType = TypeVar("FrameType", bound=AnyServerFrame)
+WaitForFutureReturnType = TypeVar("WaitForFutureReturnType")
 
-async def take_connected_frame(
-    *, frames_iter: AsyncIterable[AnyServerFrame], connection_confirmation_timeout: int
-) -> ConnectedFrame | ConnectionConfirmationTimeout:
+
+async def wait_for_or_none(
+    awaitable: Awaitable[WaitForFutureReturnType], timeout: float
+) -> WaitForFutureReturnType | None:
+    try:
+        return await asyncio.wait_for(awaitable, timeout=timeout)
+    except TimeoutError:
+        return None
+
+
+WaitForOrNone = Callable[[Awaitable[WaitForFutureReturnType], float], Awaitable[WaitForFutureReturnType | None]]
+
+
+async def take_frame_of_type(
+    *,
+    frame_type: type[FrameType],
+    frames_iter: AsyncIterable[AnyServerFrame],
+    timeout: int,
+    wait_for_or_none: WaitForOrNone[FrameType],
+) -> FrameType | list[Any]:
     collected_frames = []
 
-    async def take_connected_frame_and_collect_other_frames() -> ConnectedFrame:
+    async def inner() -> FrameType:
         async for frame in frames_iter:
-            if isinstance(frame, ConnectedFrame):
+            if isinstance(frame, frame_type):
                 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)
+    return await wait_for_or_none(inner(), timeout) or collected_frames
 
 
 def check_stomp_protocol_version(
@@ -59,18 +72,6 @@ def calculate_heartbeat_interval(*, connected_frame: ConnectedFrame, client_hear
     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: ...
@@ -87,6 +88,7 @@ class ConnectionLifespan(AbstractConnectionLifespan):
     active_subscriptions: ActiveSubscriptions
     active_transactions: ActiveTransactions
     set_heartbeat_interval: Callable[[float], None]
+    _generate_receipt_id: Callable[[], str] = field(default=lambda: _make_receipt_id())  # noqa: PLW0108
 
     async def _establish_connection(self) -> StompProtocolConnectionIssue | None:
         await self.connection.write_frame(
@@ -100,13 +102,17 @@ async def _establish_connection(self) -> StompProtocolConnectionIssue | None:
                 },
             )
         )
-        connected_frame_or_error = await take_connected_frame(
+        connected_frame_or_collected_frames = await take_frame_of_type(
+            frame_type=ConnectedFrame,
             frames_iter=self.connection.read_frames(),
-            connection_confirmation_timeout=self.connection_confirmation_timeout,
+            timeout=self.connection_confirmation_timeout,
+            wait_for_or_none=wait_for_or_none,
         )
-        if isinstance(connected_frame_or_error, ConnectionConfirmationTimeout):
-            return connected_frame_or_error
-        connected_frame = connected_frame_or_error
+        if not isinstance(connected_frame_or_collected_frames, ConnectedFrame):
+            return ConnectionConfirmationTimeout(
+                timeout=self.connection_confirmation_timeout, frames=connected_frame_or_collected_frames
+            )
+        connected_frame = connected_frame_or_collected_frames
 
         if unsupported_protocol_version_error := check_stomp_protocol_version(
             connected_frame=connected_frame, supported_version=self.protocol_version
@@ -130,10 +136,12 @@ async def enter(self) -> StompProtocolConnectionIssue | None:
 
     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()}))
-        await wait_for_receipt_frame(
+        await self.connection.write_frame(DisconnectFrame(headers={"receipt": self._generate_receipt_id()}))
+        await take_frame_of_type(
+            frame_type=ReceiptFrame,
             frames_iter=self.connection.read_frames(),
-            disconnect_confirmation_timeout=self.disconnect_confirmation_timeout,
+            timeout=self.disconnect_confirmation_timeout,
+            wait_for_or_none=wait_for_or_none,
         )
 
 

stompman/connection_manager.py

@@ -1,8 +1,7 @@
 import asyncio
 from collections.abc import AsyncGenerator, Awaitable, Callable
 from dataclasses import dataclass, field
-from types import TracebackType
-from typing import TYPE_CHECKING, Self
+from typing import TYPE_CHECKING
 
 from stompman.config import ConnectionParameters
 from stompman.connection import AbstractConnection
@@ -26,11 +25,15 @@ class ActiveConnectionState:
     lifespan: "AbstractConnectionLifespan"
 
 
+Sleep = Callable[[float], Awaitable[None]]
+
+
 async def attempt_to_connect(
     *,
     connect: Callable[[], Awaitable[ActiveConnectionState | AnyConnectionIssue]],
     connect_retry_interval: int,
     connect_retry_attempts: int,
+    sleep: Sleep,
 ) -> ActiveConnectionState:
     connection_issues = []
 
@@ -40,7 +43,7 @@ async def attempt_to_connect(
             return connection_result
 
         connection_issues.append(connection_result)
-        await asyncio.sleep(connect_retry_interval * (attempt + 1))
+        await sleep(connect_retry_interval * (attempt + 1))
 
     raise FailedAllConnectAttemptsError(retry_attempts=connect_retry_attempts, issues=connection_issues)
 
@@ -83,13 +86,10 @@ class ConnectionManager:
     _active_connection_state: ActiveConnectionState | None = field(default=None, init=False)
     _reconnect_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
 
-    async def __aenter__(self) -> Self:
+    async def enter(self) -> None:
         self._active_connection_state = await self._get_active_connection_state()
-        return self
 
-    async def __aexit__(
-        self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None
-    ) -> None:
+    async def exit(self) -> None:
         if not self._active_connection_state:
             return
         try:
@@ -131,6 +131,7 @@ async def _get_active_connection_state(self) -> ActiveConnectionState:
                 connect=self._connect_to_any_server,
                 connect_retry_interval=self.connect_retry_interval,
                 connect_retry_attempts=self.connect_retry_attempts,
+                sleep=asyncio.sleep,
             )
             return self._active_connection_state
 

tests/conftest.py

@@ -1,5 +1,5 @@
 import asyncio
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, AsyncIterable, Iterable
 from dataclasses import dataclass, field
 from typing import Any, Self, TypeVar
 
@@ -8,21 +8,13 @@
 
 import stompman
 from stompman.connection import AbstractConnection
-from stompman.connection_lifespan import AbstractConnectionLifespan
-from stompman.connection_manager import ConnectionManager
 
 
 @pytest.fixture
 def anyio_backend() -> object:
     return "asyncio"
 
 
-@pytest.fixture
-def mock_sleep(monkeypatch: pytest.MonkeyPatch) -> None:  # noqa: PT004
-    original_sleep = asyncio.sleep
-    monkeypatch.setattr("asyncio.sleep", lambda _: original_sleep(0))
-
-
 async def noop_message_handler(frame: stompman.MessageFrame) -> None: ...
 
 
@@ -52,29 +44,6 @@ class EnrichedClient(stompman.Client):
     )
 
 
-@dataclass(frozen=True, kw_only=True, slots=True)
-class NoopLifespan(AbstractConnectionLifespan):
-    connection: AbstractConnection
-    connection_parameters: stompman.ConnectionParameters
-
-    async def enter(self) -> stompman.StompProtocolConnectionIssue | None: ...
-    async def exit(self) -> None: ...
-
-
-@dataclass(kw_only=True, slots=True)
-class EnrichedConnectionManager(ConnectionManager):
-    servers: list[stompman.ConnectionParameters] = field(
-        default_factory=lambda: [stompman.ConnectionParameters("localhost", 12345, "login", "passcode")]
-    )
-    lifespan_factory: stompman.connection_lifespan.ConnectionLifespanFactory = field(default=NoopLifespan)
-    connect_retry_attempts: int = 3
-    connect_retry_interval: int = 1
-    connect_timeout: int = 3
-    read_timeout: int = 4
-    read_max_chunk_size: int = 5
-    write_retry_attempts: int = 3
-
-
 DataclassType = TypeVar("DataclassType")
 
 
@@ -140,3 +109,12 @@ def enrich_expected_frames(
         stompman.DisconnectFrame(headers={"receipt": "receipt-id-1"}),
         stompman.ReceiptFrame(headers={"receipt-id": "receipt-id-1"}),
     ]
+
+
+IterableItemT = TypeVar("IterableItemT")
+
+
+async def make_async_iter(iterable: Iterable[IterableItemT]) -> AsyncIterable[IterableItemT]:
+    for item in iterable:
+        yield item
+    await asyncio.sleep(0)

tests/test_client.py

@@ -1,29 +1,83 @@
 import asyncio
+from collections.abc import AsyncIterable, Iterable
 from unittest import mock
 
 import pytest
 
-from tests.conftest import EnrichedClient, create_spying_connection, get_read_frames_with_lifespan
+from stompman import (
+    AnyServerFrame,
+    Client,
+    ConnectedFrame,
+    ConnectFrame,
+    ConnectionParameters,
+    HeartbeatFrame,
+    ReceiptFrame,
+)
+from tests.conftest import EnrichedClient, build_dataclass, make_async_iter
 
 pytestmark = pytest.mark.anyio
 
 
-async def test_client_heartbeats_ok(monkeypatch: pytest.MonkeyPatch) -> None:
+async def test_client_heartbeats_and_lifespan_ok(monkeypatch: pytest.MonkeyPatch) -> None:
+    sleep_calls = []
+    real_sleep = asyncio.sleep
+
     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]
+    written_and_read_frames = []
 
-    async with EnrichedClient(connection_class=connection_class):
+    async def mock_read_frames(iterable: Iterable[AnyServerFrame]) -> AsyncIterable[AnyServerFrame]:
+        async for frame in make_async_iter(iterable):
+            written_and_read_frames.append(frame)
+            yield frame
+        await asyncio.Future()
+
+    connected_frame = build_dataclass(
+        ConnectedFrame, headers={"version": EnrichedClient.PROTOCOL_VERSION, "heart-beat": "1000,1000"}
+    )
+    receipt_frame = build_dataclass(ReceiptFrame)
+
+    on_heartbeat = mock.Mock()
+    write_heartbeat = mock.Mock()
+    connection = mock.AsyncMock(
+        read_frames=mock.Mock(
+            side_effect=[
+                mock_read_frames([connected_frame]),
+                mock_read_frames([HeartbeatFrame(), HeartbeatFrame()]),
+                mock_read_frames([receipt_frame]),
+            ]
+        ),
+        write_frame=mock.AsyncMock(side_effect=written_and_read_frames.append),
+        write_heartbeat=write_heartbeat,
+    )
+    connection_class = mock.Mock(connect=mock.AsyncMock(return_value=connection))
+    connection_parameters = build_dataclass(ConnectionParameters)
+    async with Client(servers=[connection_parameters], connection_class=connection_class, on_heartbeat=on_heartbeat):
         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()]
+    assert written_and_read_frames == [
+        ConnectFrame(
+            headers={
+                "accept-version": "1.2",
+                "heart-beat": "1000,1000",
+                "host": connection_parameters.host,
+                "login": connection_parameters.login,
+                "passcode": connection_parameters.unescaped_passcode,
+            }
+        ),
+        connected_frame,
+        HeartbeatFrame(),
+        HeartbeatFrame(),
+        written_and_read_frames[-2],
+        receipt_frame,
+    ]
+
+    assert sleep_calls == [0, 1, 0, 1]
+    assert write_heartbeat.mock_calls == [mock.call(), mock.call(), mock.call()]
+    assert on_heartbeat.mock_calls == [mock.call(), mock.call()]