Add ability to stop logger and add async tests

Author: thatbudakguy

mqlog/__init__.py

@@ -37,28 +37,39 @@ def __init__(
         self.capacity = capacity
         self.buffer = []
         self.will_flush = asyncio.Event()
+        self.active = True
 
-    def should_flush(self, record):
+    async def run(self):
+        """
+        Continuously publish log records via MQTT.
+        This method should be scheduled as an asyncio task.
+        """
+        self.active = True
+        while self.active:
+            await self.will_flush.wait()
+            await self._flush()
+
+    def stop(self):
+        """Stop the handler from processing further log records."""
+        self.active = False
+
+    # Check if we should publish an MQTT message
+    def _should_flush(self, record):
         return (len(self.buffer) >= self.capacity) or (
             record.levelno >= self.flush_level
         )
 
-    # Called when the logger logs a message
+    # Called by logging.Handler when the logger logs a message
     def emit(self, record):
-        self.buffer.append(record)
-        if self.should_flush(record):
+        self.buffer.append(self.format(record))
+        if self._should_flush(record):
             self.will_flush.set()
 
-    # Needs to be run in an event loop once MQTT client is connected
-    async def run(self):
-        while True:
-            await self.will_flush.wait()
-            await self._flush()
-
     # Named with an underscore to avoid conflict with logging.Handler.flush
     async def _flush(self):
         if self.buffer:
-            records = [self.buffer.pop(0) for _ in range(len(self.buffer))]
-            msg = "\n".join(self.format(record) for record in records)
-            await self.client.publish(self.topic, msg, qos=self.qos)
+            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()

tests/test_handler.py

@@ -1,31 +1,30 @@
 import asyncio
 import logging
-from unittest import IsolatedAsyncioTestCase
-from unittest.mock import AsyncMock, MagicMock, call
+from unittest import TestCase
 
 from mqlog import MqttHandler
 
 
-class TestMqttHandler(IsolatedAsyncioTestCase):
+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))
+
+
+class TestMqttHandler(TestCase):
     def setUp(self):
-        self.client = MagicMock()
-        self.client.publish = AsyncMock()
+        self.client = FakeClient()
         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)
 
-    # Run for X seconds, then stop
-    async def run_with_timeout(self, timeout=1):
-        async with asyncio.timeout(timeout):
-            await self.handler.run()
-            self.handler.close()
-
-    # Log a message at a given level after a delay of X seconds
-    async def log_message(self, msg, level, delay):
-        await asyncio.sleep(delay)
-        self.logger.log(level, msg)
-
     def test_no_flush(self):
         """Buffer should not be flushed until capacity/level reached"""
         self.handler.flush_level = logging.WARNING
@@ -45,35 +44,50 @@ def test_flush_full_buffer(self):
             self.logger.info(f"Test message {i}")
         self.assertTrue(self.handler.will_flush.is_set())
 
-    async def test_publish(self):
+    def test_publish(self):
         """Flushing should publish messages to MQTT topic"""
         self.handler.flush_level = logging.ERROR
-        await asyncio.gather(
-            self.run_with_timeout(1),
-            self.log_message("Test message 1", logging.INFO, 0.1),
-            self.log_message("Test message 2", logging.ERROR, 0.2),
-            return_exceptions=True,
-        )
-        self.client.publish.assert_awaited_with(
-            "test_topic",
-            "Test message 1\nTest message 2",
-            qos=0,
+
+        async def do_test(handler, logger):
+            logger.info("Test message 1")
+            await asyncio.sleep(0.1)
+            logger.error("Test message 2")
+            await asyncio.sleep(0.1)
+            handler.stop()
+
+        async def main(handler, logger):
+            await asyncio.gather(handler.run(), do_test(handler, logger))
+
+        asyncio.run(main(self.handler, self.logger))
+
+        self.assertEqual(
+            self.client.calls, [("test_topic", "Test message 1\nTest message 2", 0)]
         )
 
-    async def test_flush_multiple(self):
+    def test_flush_multiple(self):
         """Flushing multiple times should publish separate messages"""
         self.handler.capacity = 2
-        await asyncio.gather(
-            self.run_with_timeout(1),
-            self.log_message("Test message 1", logging.INFO, 0.1),
-            self.log_message("Test message 2", logging.INFO, 0.2),
-            self.log_message("Test message 3", logging.INFO, 0.3),
-            self.log_message("Test message 4", logging.INFO, 0.4),
-            return_exceptions=True,
-        )
-        self.client.publish.assert_has_awaits(
+
+        async def do_test(handler, logger):
+            logger.info("Test message 1")
+            await asyncio.sleep(0.1)
+            logger.info("Test message 2")
+            await asyncio.sleep(0.1)
+            logger.info("Test message 3")
+            await asyncio.sleep(0.1)
+            logger.info("Test message 4")
+            await asyncio.sleep(0.1)
+            handler.stop()
+
+        async def main(handler, logger):
+            await asyncio.gather(handler.run(), do_test(handler, logger))
+
+        asyncio.run(main(self.handler, self.logger))
+
+        self.assertEqual(
+            self.client.calls,
             [
-                call("test_topic", "Test message 1\nTest message 2", qos=0),
-                call("test_topic", "Test message 3\nTest message 4", qos=0),
-            ]
+                ("test_topic", "Test message 1\nTest message 2", 0),
+                ("test_topic", "Test message 3\nTest message 4", 0),
+            ],
         )