Heartbeat disconnected

aiostomp/aiostomp.py

@@ -393,7 +393,7 @@ def send_frame(
         headers = {} if headers is None else headers
         buf = self._protocol.build_frame(command, headers, body)
 
-        if not self._transport:
+        if not self._transport or (self.heartbeater and not self.heartbeater.connected):
             raise StompDisconnectedError()
 
         if self._stats:
@@ -448,10 +448,15 @@ async def _handle_connect(self, frame: Frame) -> None:
             sx, sy = (int(x) for x in heartbeat.split(","))
 
             if sy:
-                interval = max(self.heartbeat.get("cx", 0), sy)
-                logger.debug("Sending heartbeats every %sms", interval)
+                interval_cx = max(self.heartbeat.get("cx", 0), sy)
+                interval_cy = max(self.heartbeat.get("cy", 0), sx)
+                logger.debug("Sending heartbeats every %sms", interval_cx)
+                logger.debug("Receiving heartbeats every %sms", interval_cy)
                 self.heartbeater = StompHeartbeater(
-                    self._transport, interval=interval, loop=self._loop
+                    self._transport,
+                    interval_cx=interval_cx,
+                    interval_cy=interval_cy,
+                    loop=self._loop
                 )
                 await self.heartbeater.start()
 
@@ -490,6 +495,9 @@ def data_received(self, data: Optional[bytes]) -> None:
         if not data:
             return
 
+        if self.heartbeater:
+            self.heartbeater.receive()
+
         self._protocol.feed_data(data)
 
         for frame in self._protocol.pop_frames():

aiostomp/heartbeat.py

@@ -2,6 +2,7 @@
 from typing import Optional
 
 from contextlib import suppress
+from datetime import datetime
 
 
 class StompHeartbeater:
@@ -12,15 +13,17 @@ def __init__(
         self,
         transport: asyncio.Transport,
         loop: asyncio.AbstractEventLoop,
-        interval: int = 1000,
+        interval_cx: int = 1000,
+        interval_cy: int = 1000,
     ):
         self._transport = transport
-        self.interval = interval / 1000.0
+        self.interval_cx = interval_cx / 1000.0
+        self.interval_cy = interval_cy / 1000.0
         self.loop = loop
         self.task: Optional[asyncio.Future[None]] = None
         self.is_started = False
 
-        self.received_heartbeat = None
+        self.received_heartbeat = datetime.now()
 
     async def start(self) -> None:
         if self.is_started:
@@ -45,7 +48,15 @@ def shutdown(self) -> None:
     async def run(self) -> None:
         while True:
             await self.send()
-            await asyncio.sleep(self.interval, loop=self.loop)
+            await asyncio.sleep(self.interval_cx, loop=self.loop)
+
+    @property
+    def connected(self) -> bool:
+        elapsed = datetime.now() - self.received_heartbeat
+        return elapsed.total_seconds() < self.interval_cy * 2
 
     async def send(self) -> None:
         self._transport.write(self.HEART_BEAT)
+
+    def receive(self) -> None:
+        self.received_heartbeat = datetime.now()

tests/test_heartbeat.py

@@ -10,7 +10,7 @@ class TestStompHeartbeater(AsyncTestCase):
     async def setUpAsync(self):
         self.transport = Mock()
         self.heartbeater = StompHeartbeater(
-            self.transport, loop=asyncio.get_event_loop(), interval=100
+            self.transport, loop=asyncio.get_event_loop(), interval_cx=100, interval_cy=100
         )
 
     @patch("aiostomp.heartbeat.StompHeartbeater.stop")
@@ -74,3 +74,16 @@ async def test_can_restart_heartbeater(self, stop_mock):
 
         await asyncio.sleep(0.100)
         self.assertEqual(len(self.transport.write.call_args_list), 2)
+
+    @unittest_run_loop
+    async def test_can_monitor_connection_heartbeater(self):
+        await self.heartbeater.start()
+
+        await asyncio.sleep(0.001)
+        self.assertTrue(self.heartbeater.connected)
+
+        await asyncio.sleep(0.200)
+        self.assertFalse(self.heartbeater.connected)
+
+        self.heartbeater.receive()
+        self.assertTrue(self.heartbeater.connected)

tests/test_main.py

@@ -236,7 +236,7 @@ async def test_can_handle_connected_frame_with_heartbeat(
         await stomp._handle_connect(frame)
 
         heartbeater_klass_mock.assert_called_with(
-            stomp._transport, interval=1000, loop=self.loop
+            stomp._transport, interval_cx=1000, interval_cy=1000, loop=self.loop
         )
         heartbeater_mock.start.assert_called_once()