Provide logger class patching as a context manager

Author: s0undt3ch

pytest_reportportal/listener.py

@@ -8,7 +8,7 @@
     # This try/except can go away once we support pytest >= 3.3
     import _pytest.logging
     PYTEST_HAS_LOGGING_PLUGIN = True
-    from .rp_logging import RPLogHandler, patch_logger_class, unpatch_logger_class
+    from .rp_logging import RPLogHandler, patching_logger_class
 except ImportError:
     PYTEST_HAS_LOGGING_PLUGIN = False
 
@@ -17,6 +17,7 @@ class RPReportListener(object):
     def __init__(self, log_level=logging.NOTSET):
         # Test Item result
         self.result = None
+        self._log_level = log_level
         if PYTEST_HAS_LOGGING_PLUGIN:
             self._log_handler = RPLogHandler(log_level, filter_reportportal_client_logs=True)
 
@@ -25,13 +26,10 @@ def pytest_runtest_protocol(self, item):
         PyTestService.start_pytest_item(item)
         if PYTEST_HAS_LOGGING_PLUGIN:
             # This check can go away once we support pytest >= 3.3
-            try:
-                patch_logger_class()
+            with patching_logger_class():
                 with _pytest.logging.catching_logs(self._log_handler,
                                                    level=self._log_level):
                     yield
-            finally:
-                unpatch_logger_class()
         else:
             yield
         PyTestService.finish_pytest_item(self.result or 'SKIPPED')

pytest_reportportal/rp_logging.py

@@ -1,5 +1,6 @@
 import sys
 import logging
+from contextlib import contextmanager
 from functools import wraps
 
 from .service import PyTestService
@@ -84,51 +85,50 @@ def emit(self, record):
         )
 
 
-def patch_logger_class():
+@contextmanager
+def patching_logger_class():
     logger_class = logging.getLoggerClass()
-
-    def wrap_log(original_func):
-        @wraps(original_func)
-        def _log(self, *args, **kwargs):
-            attachment = kwargs.pop('attachment', None)
-            if attachment is not None:
-                kwargs.setdefault('extra', {}).update({'attachment': attachment})
-            return original_func(self, *args, **kwargs)
-        return _log
-
-    def wrap_makeRecord(original_func):
-        @wraps(original_func)
-        def makeRecord(self, name, level, fn, lno, msg, args, exc_info,
-                       func=None, extra=None, sinfo=None):
-            if extra is not None:
-                attachment = extra.pop('attachment', None)
-            else:
-                attachment = None
-            try:
-                # Python 3.5
-                record = original_func(self, name, level, fn, lno, msg, args,
-                                       exc_info, func=func, extra=extra,
-                                       sinfo=sinfo)
-            except TypeError:
-                # Python 2.7
-                record = original_func(self, name, level, fn, lno, msg, args,
-                                       exc_info, func=func, extra=extra)
-            record.attachment = attachment
-            return record
-        return makeRecord
-
-    # Store references to the original methods to allow unpatching
-    setattr(logger_class, '_original__log', logger_class._log)
-    logger_class._log = wrap_log(logger_class._log)
-    setattr(logger_class, '_original_makeRecord', logger_class.makeRecord)
-    logger_class.makeRecord = wrap_makeRecord(logger_class.makeRecord)
-
-
-def unpatch_logger_class():
-    logger_class = logging.getLoggerClass()
-    if hasattr(logger_class, '_original__log'):
-        logger_class._log = logger_class._original__log
-        delattr(logger_class, '_original__log')
-    if hasattr(logger_class, '_original_makeRecord'):
-        logger_class.makeRecord = logger_class._original_makeRecord
-        delattr(logger_class, '_original_makeRecord')
+    original_log = logger_class._log
+    original_makeRecord = logger_class.makeRecord
+
+    try:
+        def wrap_log(original_func):
+            @wraps(original_func)
+            def _log(self, *args, **kwargs):
+                attachment = kwargs.pop('attachment', None)
+                if attachment is not None:
+                    kwargs.setdefault('extra', {}).update(
+                        {'attachment': attachment})
+                return original_func(self, *args, **kwargs)
+            return _log
+
+        def wrap_makeRecord(original_func):
+            @wraps(original_func)
+            def makeRecord(self, name, level, fn, lno, msg, args, exc_info,
+                           func=None, extra=None, sinfo=None):
+                if extra is not None:
+                    attachment = extra.pop('attachment', None)
+                else:
+                    attachment = None
+                try:
+                    # Python 3.5
+                    record = original_func(self, name, level, fn, lno, msg,
+                                           args, exc_info, func=func,
+                                           extra=extra, sinfo=sinfo)
+                except TypeError:
+                    # Python 2.7
+                    record = original_func(self, name, level, fn, lno, msg,
+                                           args, exc_info, func=func,
+                                           extra=extra)
+                record.attachment = attachment
+                return record
+            return makeRecord
+
+        logger_class._log = wrap_log(logger_class._log)
+        logger_class.makeRecord = wrap_makeRecord(logger_class.makeRecord)
+
+        yield
+
+    finally:
+        logger_class._log = original_log
+        logger_class.makeRecord = original_makeRecord