Removed on_heartbeat callback handling from Client and tests

Author: vrslev

packages/stompman/stompman/client.py

@@ -32,7 +32,6 @@ class Client:
 
     servers: list[ConnectionParameters] = field(kw_only=False)
     on_error_frame: Callable[[ErrorFrame], Any] | None = None
-    on_heartbeat: Callable[[], Any] | Callable[[], Awaitable[Any]] | None = None
 
     heartbeat: Heartbeat = field(default=Heartbeat(1000, 1000))
     ssl: Literal[True] | SSLContext | None = None
@@ -57,7 +56,6 @@ class Client:
     _check_server_heartbeat_task: asyncio.Task[None] = field(init=False)
     _listen_task: asyncio.Task[None] = field(init=False)
     _task_group: asyncio.TaskGroup = field(init=False)
-    _on_heartbeat_is_async: bool = field(init=False)
 
     def __post_init__(self) -> None:
         self._connection_manager = ConnectionManager(
@@ -82,7 +80,6 @@ def __post_init__(self) -> None:
             check_server_alive_interval_factor=self.check_server_alive_interval_factor,
             ssl=self.ssl,
         )
-        self._on_heartbeat_is_async = inspect.iscoroutinefunction(self.on_heartbeat) if self.on_heartbeat else False
 
     async def __aenter__(self) -> Self:
         self._task_group = await self._exit_stack.enter_async_context(asyncio.TaskGroup())
@@ -157,14 +154,7 @@ async def _listen_to_frames(self) -> None:
                     case ErrorFrame():
                         if self.on_error_frame:
                             self.on_error_frame(frame)
-                    case HeartbeatFrame():
-                        if self.on_heartbeat is None:
-                            pass
-                        elif self._on_heartbeat_is_async:
-                            task_group.create_task(self.on_heartbeat())  # type: ignore[arg-type]
-                        else:
-                            self.on_heartbeat()
-                    case ConnectedFrame() | ReceiptFrame():
+                    case HeartbeatFrame() | ConnectedFrame() | ReceiptFrame():
                         pass
 
     async def send(

packages/stompman/test_stompman/test_connection_lifespan.py

@@ -1,6 +1,5 @@
 import asyncio
 from collections.abc import AsyncGenerator, Coroutine
-from functools import partial
 from typing import Any
 from unittest import mock
 
@@ -17,7 +16,6 @@
     DisconnectFrame,
     ErrorFrame,
     FailedAllConnectAttemptsError,
-    HeartbeatFrame,
     ReceiptFrame,
     UnsupportedProtocolVersion,
 )
@@ -156,58 +154,6 @@ async def mock_sleep(delay: float) -> None:
     assert write_heartbeat_mock.mock_calls == [mock.call(), mock.call(), mock.call()]
 
 
-async def test_client_on_heartbeat_none(monkeypatch: pytest.MonkeyPatch) -> None:
-    real_sleep = asyncio.sleep
-    monkeypatch.setattr("asyncio.sleep", partial(asyncio.sleep, 0))
-    connection_class, _ = create_spying_connection(
-        *get_read_frames_with_lifespan(
-            [build_dataclass(HeartbeatFrame), build_dataclass(HeartbeatFrame), build_dataclass(HeartbeatFrame)]
-        )
-    )
-
-    async with EnrichedClient(connection_class=connection_class, on_heartbeat=None):
-        await real_sleep(0)
-        await real_sleep(0)
-        await real_sleep(0)
-
-
-async def test_client_on_heartbeat_sync(monkeypatch: pytest.MonkeyPatch) -> None:
-    real_sleep = asyncio.sleep
-    monkeypatch.setattr("asyncio.sleep", partial(asyncio.sleep, 0))
-    connection_class, _ = create_spying_connection(
-        *get_read_frames_with_lifespan(
-            [build_dataclass(HeartbeatFrame), build_dataclass(HeartbeatFrame), build_dataclass(HeartbeatFrame)]
-        )
-    )
-    on_heartbeat_mock = mock.Mock()
-
-    async with EnrichedClient(connection_class=connection_class, on_heartbeat=on_heartbeat_mock):
-        await real_sleep(0)
-        await real_sleep(0)
-        await real_sleep(0)
-
-    assert on_heartbeat_mock.mock_calls == [mock.call(), mock.call(), mock.call()]
-
-
-async def test_client_on_heartbeat_async(monkeypatch: pytest.MonkeyPatch) -> None:
-    real_sleep = asyncio.sleep
-    monkeypatch.setattr("asyncio.sleep", partial(asyncio.sleep, 0))
-    connection_class, _ = create_spying_connection(
-        *get_read_frames_with_lifespan(
-            [build_dataclass(HeartbeatFrame), build_dataclass(HeartbeatFrame), build_dataclass(HeartbeatFrame)]
-        )
-    )
-    on_heartbeat_mock = mock.AsyncMock()
-
-    async with EnrichedClient(connection_class=connection_class, on_heartbeat=on_heartbeat_mock):
-        await real_sleep(0)
-        await real_sleep(0)
-        await real_sleep(0)
-
-    assert on_heartbeat_mock.await_count == 3  # noqa: PLR2004
-    assert on_heartbeat_mock.mock_calls == [mock.call.__bool__(), mock.call(), mock.call(), mock.call()]
-
-
 def test_make_receipt_id(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.undo()
     stompman.connection_lifespan._make_receipt_id()

packages/stompman/test_stompman/test_subscription.py

@@ -165,9 +165,7 @@ async def test_client_listen_routing_ok(monkeypatch: pytest.MonkeyPatch, faker:
     second_message_handler, second_error_handler = mock.AsyncMock(side_effect=SomeError), mock.Mock()
 
     async with EnrichedClient(
-        connection_class=connection_class,
-        on_error_frame=(on_error_frame := mock.Mock()),
-        on_heartbeat=(on_heartbeat := mock.Mock()),
+        connection_class=connection_class, on_error_frame=(on_error_frame := mock.Mock())
     ) as client:
         first_subscription = await client.subscribe(
             faker.pystr(), handler=first_message_handler, on_suppressed_exception=first_error_handler
@@ -187,7 +185,6 @@ async def test_client_listen_routing_ok(monkeypatch: pytest.MonkeyPatch, faker:
     second_error_handler.assert_called_once_with(SomeError(), third_message_frame)
 
     on_error_frame.assert_called_once_with(error_frame)
-    on_heartbeat.assert_called_once_with()
 
 
 @pytest.mark.parametrize("side_effect", [None, SomeError])