Adjustments for heartbeat logging and connection manager logic

Author: vrslev

packages/stompman/stompman/client.py

@@ -109,6 +109,7 @@ async def __aexit__(
     def _restart_heartbeat_tasks(self, server_heartbeat: Heartbeat) -> None:
         self._send_heartbeat_task.cancel()
         self._check_server_heartbeat_task.cancel()
+        print(server_heartbeat)
         self._send_heartbeat_task = self._task_group.create_task(
             self._send_heartbeats_forever(server_heartbeat.want_to_receive_interval_ms)
         )

packages/stompman/stompman/connection_manager.py

@@ -8,7 +8,6 @@
 
 from stompman.config import ConnectionParameters, Heartbeat
 from stompman.connection import AbstractConnection
-from stompman.connection_lifespan import EstablishedConnectionResult
 from stompman.errors import (
     AllServersUnavailable,
     AnyConnectionIssue,
@@ -85,6 +84,8 @@ async def _create_connection_to_any_server(self) -> tuple[AbstractConnection, Co
         return None
 
     async def _connect_to_any_server(self) -> ActiveConnectionState | AnyConnectionIssue:
+        from stompman.connection_lifespan import EstablishedConnectionResult  # noqa: PLC0415
+
         if not (connection_and_server := await self._create_connection_to_any_server()):
             return AllServersUnavailable(servers=self.servers, timeout=self.connect_timeout)
         connection, connection_parameters = connection_and_server

packages/stompman/test_stompman/conftest.py

@@ -125,7 +125,9 @@ async def read_frames() -> AsyncGenerator[stompman.AnyServerFrame, None]:
         "passcode": "passcode",
     },
 )
-CONNECTED_FRAME = stompman.ConnectedFrame(headers={"version": stompman.Client.PROTOCOL_VERSION, "heart-beat": "1,1"})
+CONNECTED_FRAME = stompman.ConnectedFrame(
+    headers={"version": stompman.Client.PROTOCOL_VERSION, "heart-beat": "1000,1000"}
+)
 
 
 @pytest.fixture(autouse=True)

packages/stompman/test_stompman/test_connection_lifespan.py

@@ -152,7 +152,7 @@ async def mock_sleep(delay: float) -> None:
         await real_sleep(0)
         await real_sleep(0)
 
-    assert sleep_calls == [0, 1, 1]
+    assert sleep_calls == [0, 0, 1, 3, 1, 3]
     assert write_heartbeat_mock.mock_calls == [mock.call(), mock.call(), mock.call()]
 
 

packages/stompman/test_stompman/test_connection_manager.py

@@ -132,7 +132,7 @@ class MockConnection(BaseMockConnection):
 
 
 async def test_get_active_connection_state_lifespan_flaky_ok() -> None:
-    enter = mock.AsyncMock(side_effect=[ConnectionLostError, None])
+    enter = mock.AsyncMock(side_effect=[ConnectionLostError, build_dataclass(EstablishedConnectionResult)])
     lifespan_factory = mock.Mock(return_value=mock.Mock(enter=enter))
     manager = EnrichedConnectionManager(lifespan_factory=lifespan_factory, connection_class=BaseMockConnection)
 
@@ -206,7 +206,8 @@ async def test_connection_manager_context_lifespan_aexit_raises_connection_lost(
     async with EnrichedConnectionManager(
         lifespan_factory=mock.Mock(
             return_value=mock.Mock(
-                enter=mock.AsyncMock(return_value=None), exit=mock.AsyncMock(side_effect=[ConnectionLostError])
+                enter=mock.AsyncMock(return_value=build_dataclass(EstablishedConnectionResult)),
+                exit=mock.AsyncMock(side_effect=[ConnectionLostError]),
             )
         ),
         connection_class=BaseMockConnection,
@@ -222,7 +223,11 @@ class MockConnection(BaseMockConnection):
         close = connection_close
 
     async with EnrichedConnectionManager(
-        lifespan_factory=mock.Mock(return_value=mock.Mock(enter=mock.AsyncMock(return_value=None), exit=close)),
+        lifespan_factory=mock.Mock(
+            return_value=mock.Mock(
+                enter=mock.AsyncMock(return_value=build_dataclass(EstablishedConnectionResult)), exit=close
+            )
+        ),
         connection_class=MockConnection,
     ):
         pass