Ensure buffer can't overflow; update uptime check

This takes advantage of amqc's MQTTClient.up Event flag to wait
until the connection is good before trying to flush.

Also truncates the buffer while the connection is down so that
it doesn't grow indefinitely.

Author: thatbudakguy

mqlog/__init__.py

@@ -47,6 +47,7 @@ async def run(self):
         This method should be scheduled as an asyncio task.
         """
         while True:
+            await self.client.up.wait()
             await self.will_flush.wait()
             await self._flush()
 
@@ -59,6 +60,11 @@ def _should_flush(self, record):
     # Called by logging.Handler when the logger logs a message
     def emit(self, record):
         self.buffer.append(self.format(record))
+
+        # Prevent buffer from growing indefinitely
+        if self.buffer and len(self.buffer) > self.capacity:
+            self.buffer = self.buffer[-self.capacity :]
+
         if self._should_flush(record):
             self.will_flush.set()
 
@@ -73,15 +79,11 @@ async def _flush(self):
             # Try to publish, but if we fail, just log the error – don't
             # want to cause cascading errors. Something else is responsible
             # for bringing the connection back up.
+            msg = "\n".join([line for line in self.buffer])
             try:
-                msg = "\n".join([line for line in self.buffer])
                 await self.client.publish(self.topic, msg, qos=self.qos)
                 self.buffer.clear()
             except Exception as e:
                 self._logger.error(f"Failed to publish logs via MQTT: {e}")
-            self.will_flush.clear()
-
-        # If we were supposed to flush but couldn't, truncate the buffer
-        # to prevent it from growing indefinitely
-        if self.buffer and len(self.buffer) >= self.capacity:
-            self.buffer = self.buffer[-self.capacity :]
+            finally:
+                self.will_flush.clear()

tests/test_handler.py

@@ -1,6 +1,7 @@
 import asyncio
 import logging
-from unittest import TestCase
+import unittest
+from unittest import TestCase, skip
 
 from mqlog import MqttHandler
 from tests.utils import AsyncMock, Mock, call
@@ -9,23 +10,25 @@
 class FakeClient:
     """A fake MQTT client for testing purposes."""
 
-    publish: None
+    def __init__(self):
+        self.publish = AsyncMock()
+        self.up = asyncio.Event()
 
 
 class TestMqttHandler(TestCase):
     def setUp(self):
         self.client = FakeClient()
-        self.client.publish = AsyncMock()
         self.handler = MqttHandler(self.client, "test_topic")
-        self.handler.setFormatter(logging.Formatter("%(message)s"))
         self.logger = logging.getLogger("test")
         self.logger.addHandler(self.handler)
         self.logger.setLevel(logging.INFO)
+        self.client.up.set()  # Simulate client being ready
 
     def tearDown(self):
         """Clean up after each test"""
         self.handler.buffer.clear()
         self.handler.will_flush.clear()
+        self.client.up.clear()
 
     def test_default_formatter(self):
         """Handler should use default formatter if none is set"""
@@ -54,7 +57,7 @@ def test_flush_full_buffer(self):
     def test_flush_fail(self):
         """Flushing should log an error if publish fails"""
         self.handler._logger.error = Mock()
-        self.client.publish = AsyncMock(side_effect=Exception("Publish failed"))
+        self.client.publish.side_effect = Exception("Publish failed")
 
         async def do_test(handler: logging.Handler, logger: logging.Logger):
             asyncio.create_task(handler.run())
@@ -81,7 +84,7 @@ async def do_test(handler: logging.Handler, logger: logging.Logger):
         asyncio.run(do_test(self.handler, self.logger))
 
         self.client.publish.assert_called_with(
-            "test_topic", "Test message 1\nTest message 2", qos=0
+            "test_topic", "INFO:test:Test message 1\nERROR:test:Test message 2", qos=0
         )
 
     def test_flush_multiple(self):
@@ -104,15 +107,23 @@ async def do_test(handler: logging.Handler, logger: logging.Logger):
 
         self.client.publish.assert_has_calls(
             [
-                call("test_topic", "Test message 1\nTest message 2", qos=0),
-                call("test_topic", "Test message 3\nTest message 4", qos=0),
+                call(
+                    "test_topic",
+                    "INFO:test:Test message 1\nINFO:test:Test message 2",
+                    qos=0,
+                ),
+                call(
+                    "test_topic",
+                    "INFO:test:Test message 3\nINFO:test:Test message 4",
+                    qos=0,
+                ),
             ]
         )
 
     def test_buffer_overflow(self):
         """Buffer should get truncated if it exceeds capacity"""
         self.handler.capacity = 3
-        self.client.publish = AsyncMock(side_effect=Exception("Publish failed"))
+        self.client.up.clear()  # Simulate client being disconnected
 
         async def do_test(handler: logging.Handler, logger: logging.Logger):
             asyncio.create_task(handler.run())
@@ -129,5 +140,36 @@ async def do_test(handler: logging.Handler, logger: logging.Logger):
         asyncio.run(do_test(self.handler, self.logger))
 
         self.assertEqual(
-            self.handler.buffer, ["Test message 2", "Test message 3", "Test message 4"]
+            self.handler.buffer,
+            [
+                "INFO:test:Test message 2",
+                "INFO:test:Test message 3",
+                "INFO:test:Test message 4",
+            ],
+        )
+
+    def test_reconnect(self):
+        """Handler should reconnect and publish messages after a disconnect"""
+        self.client.up.clear()  # Simulate client being disconnected
+
+        async def do_test(
+            handler: logging.Handler, logger: logging.Logger, client: FakeClient
+        ):
+            asyncio.create_task(handler.run())
+            await asyncio.sleep(0.1)
+            logger.error("Test message 1")
+            await asyncio.sleep(0.1)
+
+            # Simulate reconnect
+            client.up.set()
+            await asyncio.sleep(0.1)
+
+        asyncio.run(do_test(self.handler, self.logger, self.client))
+
+        self.client.publish.assert_called_with(
+            "test_topic", "ERROR:test:Test message 1", qos=0
         )
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/utils.py

@@ -6,16 +6,16 @@ class Mock:
     """A mock callable object that stores its calls."""
 
     def __init__(self, return_value=None, side_effect=None):
-        self._return_value = return_value
-        self._side_effect = side_effect
+        self.return_value = return_value
+        self.side_effect = side_effect
         self._calls = []
 
     def __call__(self, *args, **kwargs):
         """Call the mock and store the call details."""
         self._calls.append(call(*args[1:], **kwargs))  # Skip the first argument (self)
-        if self._side_effect:
-            raise self._side_effect
-        return self._return_value
+        if self.side_effect:
+            raise self.side_effect
+        return self.return_value
 
     def assert_called(self):
         """Assert that the mock was called at least once."""