logrecorder.py

class LogRecorder():
    '''Context manager which captures log records for testing.

    Example usage in a TestCase:

        def capture_logs(self, logger):
            return LogRecorder(logger)

        def assertLogged(self, recorder, message):
            self.assertIn(message, recorder)

        def test_something(self):
            # Do something which generates log messages
            with self.capture_logs(logger) as l:
                self.assertLogged(l, 'my message')
                
            # Or
            with LogRecorder(logger) as l:
                self.assertIn('my  message', l)
    '''

    class LogHandler(logging.Handler):
        def __init__(self, recorder):
            self.recorder = recorder
            super().__init__()

        def emit(self, record):
            self.recorder.logs.append(record)

    def __init__(self, logger):
        self.logger = logger
        self.handler = self.LogHandler(self)
        self.logs = []

    def __enter__(self):
        self.logger.addHandler(self.handler)
        return self

    def __exit__(self, type, value, traceback):
        self.logger.removeHandler(self.handler)

    def __contains__(self, message):
        messages = (l.getMessage() for l in self.logs)
        return any(message in m for m in messages)