tests: Add a test case that verifies behavior when the log record factory is modified

The test currently fails with the following error:

  $ PYTHONPATH=$(pwd)/tests:$PYTHONPATH python3 -m unittest -v log_handler_test
  ...
  test_logrecordfactory (log_handler_test.LogRecordFactoryTest.test_logrecordfactory) ... --- Logging error ---
  Traceback (most recent call last):
    File "/usr/lib/python3.12/logging/__init__.py", line 464, in format
      return self._format(record)
             ^^^^^^^^^^^^^^^^^^^^
    File "/usr/lib/python3.12/logging/__init__.py", line 460, in _format
      return self._fmt % values
             ~~~~~~~~~~^~~~~~~~
  KeyError: 'custom_attribute'

  During handling of the above exception, another exception occurred:

  Traceback (most recent call last):
    File "/usr/lib/python3.12/logging/__init__.py", line 1160, in emit
      msg = self.format(record)
            ^^^^^^^^^^^^^^^^^^^
    File "/usr/lib/python3.12/logging/__init__.py", line 999, in format
      return fmt.format(record)
             ^^^^^^^^^^^^^^^^^^
    File "/usr/lib/python3.12/logging/__init__.py", line 999, in format
      return fmt.format(record)
             ^^^^^^^^^^^^^^^^^^
    File "/usr/lib/python3.12/logging/__init__.py", line 706, in format
      s = self.formatMessage(record)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/usr/lib/python3.12/logging/__init__.py", line 675, in formatMessage
      return self._style.format(record)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/usr/lib/python3.12/logging/__init__.py", line 466, in format
      raise ValueError('Formatting field not found in record: %s' % e)
  ValueError: Formatting field not found in record: 'custom_attribute'

Signed-off-by: Marc Hartmayer <[email protected]>

Author: mhartmay

tests/log_handler_test.py

@@ -14,11 +14,28 @@
 import mitogen.utils
 from mitogen.core import b
 
+PY2 = sys.version_info[0] == 2
+if PY2:
+
+    def logging_getLogRecordFactory():
+        return logging.LogRecord
+
+    def logging_setLogRecordFactory(factory):
+        logging.LogRecord = factory
+
+else:
+    logging_getLogRecordFactory = logging.getLogRecordFactory
+    logging_setLogRecordFactory = logging.setLogRecordFactory
+
 
 def ping():
     pass
 
 
+def log_test():
+    logging.getLogger(__name__).info("This is a test")
+
+
 class BufferingTest(testlib.TestCase):
     klass = mitogen.core.LogHandler
 
@@ -89,6 +106,36 @@ def test_earliest_messages_logged_via(self):
         expect = 'Parent is context %s (%s)' % (c1.context_id, 'parent')
         self.assertIn(expect, logs)
 
+
+class LogRecordFactoryTest(testlib.RouterMixin, testlib.TestCase):
+    def setUp(self):
+        super(LogRecordFactoryTest, self).setUp()
+        self.original_factory = logging_getLogRecordFactory()
+
+    def tearDown(self):
+        logging_setLogRecordFactory(self.original_factory)
+        super(LogRecordFactoryTest, self).tearDown()
+
+    def test_logrecordfactory(self):
+        # Change logging factory and add a custom attribute
+        old_factory = logging_getLogRecordFactory()
+
+        def record_factory(*args, **kwargs):
+            record = old_factory(*args, **kwargs)
+            record.custom_attribute = 0xDEADBEEF
+            return record
+
+        logging_setLogRecordFactory(record_factory)
+        c1 = self.router.local(name="c1")
+        log = testlib.LogCapturer(
+            __name__, formatter=logging.Formatter("%(custom_attribute)x - %(message)s")
+        )
+        log.start()
+        c1.call(log_test)
+        logs = log.stop()
+        self.assertIn("deadbeef - This is a test", logs)
+
+
 StartupTest = unittest.skipIf(
     condition=sys.version_info < (2, 7) or sys.version_info >= (3, 6),
     reason="Message log flaky on Python < 2.7 or >= 3.6"

tests/testlib.py

@@ -393,10 +393,13 @@ def emit(self, msg):
 
 
 class LogCapturer(object):
-    def __init__(self, name=None):
+    def __init__(self, name=None, formatter=None):
         self.sio = StringIO()
         self.logger = logging.getLogger(name)
-        self.handler = CaptureStreamHandler(self.sio)
+        handler = CaptureStreamHandler(self.sio)
+        if formatter is not None:
+            handler.setFormatter(formatter)
+        self.handler = handler
         self.old_propagate = self.logger.propagate
         self.old_handlers = self.logger.handlers
         self.old_level = self.logger.level