Context manager for testing log output

log_context.py

import logging

from contextlib import contextmanager
from cStringIO import StringIO

@contextmanager
def log_context(logger_name, level=logging.DEBUG, formatter=None):
    """
    :param logger_name: Name of your logger as passed to ``logging.getLogger``.
    :param level: (optional) logging level to capture.
    :param formatter: (optional) logging.Formatter to use.

    Use this context manager to help test log output. Suppose you have an app 
    that is required to log some of data. We could create a test for that like
    this ::

        # unittest method
        def test_log_output(self)
            with log_context("app-logger") as log:
                self.test_client.get("/foo")
                assert "Some String" in log.getvalue()

    """
    log = logging.getLogger(logger_name)
    log.setLevel(level)
    stream = StringIO()
    handler = logging.StreamHandler(stream)

    if formatter:
        handler.setFormatter(formatter)

    try:
        log.addHandler(handler)
        yield stream
    finally:
        stream.close()
        log.removeHandler(handler)