From 20e0626e424276c83cd1f2eb42fdddb7cf56072e Mon Sep 17 00:00:00 2001
From: Jean-Paul Calderone <exarkun@twistedmatrix.com>
Date: Mon, 29 Nov 2021 13:27:17 -0500
Subject: [PATCH] add capture_logging that parameterizes JSON encoder

---
 src/allmydata/util/_eliot_updates.py | 100 ++++++++++++++++++++++++++-
 src/allmydata/util/eliotutil.py      |  13 +---
 2 files changed, 102 insertions(+), 11 deletions(-)

diff --git a/src/allmydata/util/_eliot_updates.py b/src/allmydata/util/_eliot_updates.py
index 4ff0caf4d..8e3beca45 100644
--- a/src/allmydata/util/_eliot_updates.py
+++ b/src/allmydata/util/_eliot_updates.py
@@ -19,12 +19,17 @@ if PY2:
     from builtins import filter, map, zip, ascii, chr, hex, input, next, oct, open, pow, round, super, bytes, dict, list, object, range, str, max, min  # noqa: F401
 
 import json as pyjson
-from functools import partial
+from functools import wraps, partial
 
 from eliot import (
     MemoryLogger as _MemoryLogger,
 )
 
+from eliot.testing import (
+    check_for_errors,
+    swap_logger,
+)
+
 from .jsonbytes import AnyBytesJSONEncoder
 
 # There are currently a number of log messages that include non-UTF-8 bytes.
@@ -86,3 +91,96 @@ if PY2:
     MemoryLogger = partial(_CustomEncoderMemoryLogger, encoder=eliot_json_encoder)
 else:
     MemoryLogger = partial(_MemoryLogger, encoder=eliot_json_encoder)
+
+def validateLogging(
+    assertion, *assertionArgs, **assertionKwargs
+):
+    """
+    Decorator factory for L{unittest.TestCase} methods to add logging
+    validation.
+
+    1. The decorated test method gets a C{logger} keyword argument, a
+       L{MemoryLogger}.
+    2. All messages logged to this logger will be validated at the end of
+       the test.
+    3. Any unflushed logged tracebacks will cause the test to fail.
+
+    For example:
+
+        from unittest import TestCase
+        from eliot.testing import assertContainsFields, validateLogging
+
+        class MyTests(TestCase):
+            def assertFooLogging(self, logger):
+                assertContainsFields(self, logger.messages[0], {"key": 123})
+
+
+    @param assertion: A callable that will be called with the
+       L{unittest.TestCase} instance, the logger and C{assertionArgs} and
+       C{assertionKwargs} once the actual test has run, allowing for extra
+       logging-related assertions on the effects of the test. Use L{None} if you
+       want the cleanup assertions registered but no custom assertions.
+
+    @param assertionArgs: Additional positional arguments to pass to
+        C{assertion}.
+
+    @param assertionKwargs: Additional keyword arguments to pass to
+        C{assertion}.
+
+    @param encoder_: C{json.JSONEncoder} subclass to use when validating JSON.
+    """
+    encoder_ = assertionKwargs.pop("encoder_", eliot_json_encoder)
+    def decorator(function):
+        @wraps(function)
+        def wrapper(self, *args, **kwargs):
+            skipped = False
+
+            kwargs["logger"] = logger = MemoryLogger(encoder=encoder_)
+            self.addCleanup(check_for_errors, logger)
+            # TestCase runs cleanups in reverse order, and we want this to
+            # run *before* tracebacks are checked:
+            if assertion is not None:
+                self.addCleanup(
+                    lambda: skipped
+                    or assertion(self, logger, *assertionArgs, **assertionKwargs)
+                )
+            try:
+                return function(self, *args, **kwargs)
+            except self.skipException:
+                skipped = True
+                raise
+
+        return wrapper
+
+    return decorator
+
+# PEP 8 variant:
+validate_logging = validateLogging
+
+def capture_logging(
+    assertion, *assertionArgs, **assertionKwargs
+):
+    """
+    Capture and validate all logging that doesn't specify a L{Logger}.
+
+    See L{validate_logging} for details on the rest of its behavior.
+    """
+    encoder_ = assertionKwargs.pop("encoder_", eliot_json_encoder)
+    def decorator(function):
+        @validate_logging(
+            assertion, *assertionArgs, encoder_=encoder_, **assertionKwargs
+        )
+        @wraps(function)
+        def wrapper(self, *args, **kwargs):
+            logger = kwargs["logger"]
+            previous_logger = swap_logger(logger)
+
+            def cleanup():
+                swap_logger(previous_logger)
+
+            self.addCleanup(cleanup)
+            return function(self, *args, **kwargs)
+
+        return wrapper
+
+    return decorator
diff --git a/src/allmydata/util/eliotutil.py b/src/allmydata/util/eliotutil.py
index 5067876c5..789ef38ff 100644
--- a/src/allmydata/util/eliotutil.py
+++ b/src/allmydata/util/eliotutil.py
@@ -23,6 +23,7 @@ __all__ = [
     "opt_help_eliot_destinations",
     "validateInstanceOf",
     "validateSetMembership",
+    "capture_logging",
 ]
 
 from future.utils import PY2
@@ -33,7 +34,7 @@ from six import ensure_text
 from sys import (
     stdout,
 )
-from functools import wraps, partial
+from functools import wraps
 from logging import (
     INFO,
     Handler,
@@ -67,8 +68,6 @@ from eliot.twisted import (
     DeferredContext,
     inline_callbacks,
 )
-from eliot.testing import capture_logging as eliot_capture_logging
-
 from twisted.python.usage import (
     UsageError,
 )
@@ -91,6 +90,7 @@ from twisted.application.service import Service
 from ._eliot_updates import (
     MemoryLogger,
     eliot_json_encoder,
+    capture_logging,
 )
 
 def validateInstanceOf(t):
@@ -330,10 +330,3 @@ def log_call_deferred(action_type):
                 return DeferredContext(d).addActionFinish()
         return logged_f
     return decorate_log_call_deferred
-
-# On Python 3, encoding bytes to JSON doesn't work, so we have a custom JSON
-# encoder we want to use when validating messages.
-if PY2:
-    capture_logging = eliot_capture_logging
-else:
-    capture_logging = partial(eliot_capture_logging, encoder_=AnyBytesJSONEncoder)