Log but don't raise when publish fails

Closes #4

Author: thatbudakguy

mqlog/__init__.py

@@ -39,6 +39,7 @@ def __init__(
         self.capacity = capacity
         self.buffer = []
         self.will_flush = asyncio.Event()
+        self._logger = logging.getLogger("mqlog")
 
     async def run(self):
         """
@@ -69,8 +70,10 @@ def format(self, record):
     # Named with an underscore to avoid conflict with logging.Handler.flush
     async def _flush(self):
         if self.buffer:
-            msg = ""
-            while self.buffer:
-                msg += self.buffer.pop(0) + "\n"
-            await self.client.publish(self.topic, msg.strip(), qos=self.qos)
-        self.will_flush.clear()
+            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("Failed to publish logs", exc_info=e)
+            self.will_flush.clear()

tests/test_handler.py

@@ -4,22 +4,19 @@
 from unittest import TestCase, skipIf
 
 from mqlog import MqttHandler
+from tests.utils import AsyncMock, Mock, call
 
 
 class FakeClient:
     """A fake MQTT client for testing purposes."""
 
-    def __init__(self):
-        self.calls = []
-
-    # Store the call like a mock so we can check it later
-    async def publish(self, topic, payload, qos=0):
-        self.calls.append((topic, payload, qos))
+    publish: None
 
 
 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")
@@ -50,32 +47,46 @@ def test_flush_full_buffer(self):
             self.logger.info(f"Test message {i}")
         self.assertTrue(self.handler.will_flush.is_set())
 
-    @skipIf(os.getenv("CI"), "Hangs in CI")
+    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"))
+
+        async def do_test(handler: logging.Handler, logger: logging.Logger):
+            asyncio.create_task(handler.run())
+            await asyncio.sleep(0.1)  # Allow handler to start
+            logger.error("Test message")  # should trigger flush
+            await asyncio.sleep(0.1)
+
+        asyncio.run(do_test(self.handler, self.logger))
+
+        self.handler._logger.error.assert_called()
+
     def test_publish(self):
         """Flushing should publish messages to MQTT topic"""
         self.handler.flush_level = logging.ERROR
 
-        async def do_test(logger):
+        async def do_test(handler: logging.Handler, logger: logging.Logger):
+            asyncio.create_task(handler.run())
+            await asyncio.sleep(0.1)
             logger.info("Test message 1")
             await asyncio.sleep(0.1)
             logger.error("Test message 2")
             await asyncio.sleep(0.1)
 
-        async def main(handler, logger):
-            await asyncio.gather(handler.run(), do_test(logger))
-
-        asyncio.run(main(self.handler, self.logger))
+        asyncio.run(do_test(self.handler, self.logger))
 
-        self.assertEqual(
-            self.client.calls, [("test_topic", "Test message 1\nTest message 2", 0)]
+        self.client.publish.assert_called_with(
+            "test_topic", "Test message 1\nTest message 2", qos=0
         )
 
-    @skipIf(os.getenv("CI"), "Hangs in CI")
     def test_flush_multiple(self):
         """Flushing multiple times should publish separate messages"""
         self.handler.capacity = 2
 
-        async def do_test(logger):
+        async def do_test(handler: logging.Handler, logger: logging.Logger):
+            asyncio.create_task(handler.run())
+            await asyncio.sleep(0.1)
             logger.info("Test message 1")
             await asyncio.sleep(0.1)
             logger.info("Test message 2")
@@ -85,15 +96,11 @@ async def do_test(logger):
             logger.info("Test message 4")
             await asyncio.sleep(0.1)
 
-        async def main(handler, logger):
-            await asyncio.gather(handler.run(), do_test(logger))
-
-        asyncio.run(main(self.handler, self.logger))
+        asyncio.run(do_test(self.handler, self.logger))
 
-        self.assertEqual(
-            self.client.calls,
+        self.client.publish.assert_has_calls(
             [
-                ("test_topic", "Test message 1\nTest message 2", 0),
-                ("test_topic", "Test message 3\nTest message 4", 0),
-            ],
+                call("test_topic", "Test message 1\nTest message 2", qos=0),
+                call("test_topic", "Test message 3\nTest message 4", qos=0),
+            ]
         )

tests/utils.py

@@ -1,6 +1,3 @@
-# from unittest import TestCase
-
-
 def call(*args, **kwargs):
     return (tuple(args), dict(kwargs))
 
@@ -14,7 +11,8 @@ def __init__(self, return_value=None, side_effect=None):
         self._calls = []
 
     def __call__(self, *args, **kwargs):
-        self._calls.append(call(*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
@@ -29,17 +27,22 @@ def assert_not_called(self):
 
     def assert_called_with(self, *args, **kwargs):
         """Assert that the mock was last called with the given arguments."""
-        # First call should be self, so we prepend it
-        expected_args = [self] + list(args)
-        expectation = call(*expected_args, **kwargs)
+        # Fail if no calls were made
+        self.assert_called()
 
         # Try to have a useful output for assertion failures
+        expectation = call(*args, **kwargs)
         assert self._calls[-1] == expectation, "Expected call with {}, got {}".format(
             expectation, self._calls[-1]
         )
 
     def assert_has_calls(self, calls):
         """Assert that the mock has the expected calls with arguments."""
+        assert calls, "Expected calls cannot be empty."
+
+        # Fail if no calls were made
+        self.assert_called()
+
         assert self._calls == calls, "Expected calls {}, got {}".format(
             calls, self._calls
         )