Revert tests refactoring (#59)

Author: vrslev

pyproject.toml

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

stompman/__init__.py

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

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._connection_manager.enter()
+        await self._exit_stack.enter_async_context(self._connection_manager)
         self._listen_task = self._task_group.create_task(self._listen_to_frames())
         return self
 
@@ -91,7 +91,6 @@ 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,14 +1,14 @@
 import asyncio
-from collections.abc import AsyncIterable, Awaitable, Callable
-from dataclasses import dataclass, field
-from typing import Any, Protocol, TypeVar
+from collections.abc import Callable
+from contextlib import suppress
+from dataclasses import dataclass
+from typing import Protocol
 from uuid import uuid4
 
 from stompman.config import ConnectionParameters, Heartbeat
 from stompman.connection import AbstractConnection
 from stompman.errors import ConnectionConfirmationTimeout, StompProtocolConnectionIssue, UnsupportedProtocolVersion
 from stompman.frames import (
-    AnyServerFrame,
     ConnectedFrame,
     ConnectFrame,
     DisconnectFrame,
@@ -21,56 +21,6 @@
 )
 from stompman.transaction import ActiveTransactions, commit_pending_transactions
 
-FrameType = TypeVar("FrameType", bound=AnyServerFrame)
-WaitForFutureReturnType = TypeVar("WaitForFutureReturnType")
-
-
-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 inner() -> FrameType:
-        async for frame in frames_iter:
-            if isinstance(frame, frame_type):
-                return frame
-            collected_frames.append(frame)
-        msg = "unreachable"
-        raise AssertionError(msg)
-
-    return await wait_for_or_none(inner(), timeout) or 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
-
 
 class AbstractConnectionLifespan(Protocol):
     async def enter(self) -> StompProtocolConnectionIssue | None: ...
@@ -88,7 +38,6 @@ 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(
@@ -102,54 +51,61 @@ async def _establish_connection(self) -> StompProtocolConnectionIssue | None:
                 },
             )
         )
-        connected_frame_or_collected_frames = await take_frame_of_type(
-            frame_type=ConnectedFrame,
-            frames_iter=self.connection.read_frames(),
-            timeout=self.connection_confirmation_timeout,
-            wait_for_or_none=wait_for_or_none,
-        )
-        if not isinstance(connected_frame_or_collected_frames, ConnectedFrame):
-            return ConnectionConfirmationTimeout(
-                timeout=self.connection_confirmation_timeout, frames=connected_frame_or_collected_frames
+        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
             )
-        connected_frame = connected_frame_or_collected_frames
+        except TimeoutError:
+            return ConnectionConfirmationTimeout(timeout=self.connection_confirmation_timeout, frames=collected_frames)
 
-        if unsupported_protocol_version_error := check_stomp_protocol_version(
-            connected_frame=connected_frame, supported_version=self.protocol_version
-        ):
-            return unsupported_protocol_version_error
+        if connected_frame.headers["version"] != self.protocol_version:
+            return UnsupportedProtocolVersion(
+                given_version=connected_frame.headers["version"], supported_version=self.protocol_version
+            )
 
+        server_heartbeat = Heartbeat.from_header(connected_frame.headers["heart-beat"])
         self.set_heartbeat_interval(
-            calculate_heartbeat_interval(connected_frame=connected_frame, client_heartbeat=self.client_heartbeat)
+            max(self.client_heartbeat.will_send_interval_ms, server_heartbeat.want_to_receive_interval_ms) / 1000
         )
         return None
 
     async def enter(self) -> StompProtocolConnectionIssue | None:
-        if protocol_connection_issue := await self._establish_connection():
-            return protocol_connection_issue
-
+        if connection_issue := await self._establish_connection():
+            return 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": self._generate_receipt_id()}))
-        await take_frame_of_type(
-            frame_type=ReceiptFrame,
-            frames_iter=self.connection.read_frames(),
-            timeout=self.disconnect_confirmation_timeout,
-            wait_for_or_none=wait_for_or_none,
-        )
+        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())
 
 
 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,7 +1,8 @@
 import asyncio
-from collections.abc import AsyncGenerator, Awaitable, Callable
+from collections.abc import AsyncGenerator
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING
+from types import TracebackType
+from typing import TYPE_CHECKING, Self
 
 from stompman.config import ConnectionParameters
 from stompman.connection import AbstractConnection
@@ -25,52 +26,6 @@ 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 = []
-
-    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 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]
@@ -86,10 +41,13 @@ class ConnectionManager:
     _active_connection_state: ActiveConnectionState | None = field(default=None, init=False)
     _reconnect_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
 
-    async def enter(self) -> None:
+    async def __aenter__(self) -> Self:
         self._active_connection_state = await self._get_active_connection_state()
+        return self
 
-    async def exit(self) -> None:
+    async def __aexit__(
+        self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None
+    ) -> None:
         if not self._active_connection_state:
             return
         try:
@@ -112,28 +70,47 @@ async def _create_connection_to_one_server(self, server: ConnectionParameters) -
             )
         return None
 
-    async def _connect_to_any_server(self) -> ActiveConnectionState | AnyConnectionIssue:
-        active_connection_state = await connect_to_first_server(
+    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]
-        )
-        return await make_healthy_connection(
-            active_connection_state=active_connection_state, servers=self.servers, connect_timeout=self.connect_timeout
-        )
+        ):
+            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
 
     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
-            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,
-                sleep=asyncio.sleep,
-            )
-            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)
 
     def _clear_active_connection_state(self) -> None:
         self._active_connection_state = None