gistfile1.rst

Testing standard streams in Python

I was testing a program that took N lines of input, manipulated them somehow and returned M lines of output. To test it properly I wanted to test both the internal interface and the external interface.

The main function looked something like this:

import sys

def main():
    for line in sys.stdin:
        print do_something(line)

The internal interface, the do_something function was easy to test. The test just called it with an example line of input and compared the output with expected string.

Testing the external interface was a bit harder, because I had to capture sys.stdout and provide my IO object for sys.stdin. One way to do it is to override setUp and tearDown methods, but I didn't want to capture standard IO streams for all the test cases in the class. Instead I used contextmanager and it made my tests look beautiful. Here's the test code:

from StringIO import StringIO

def test_main():
    input = StringIO('foo')

    with provided_stdin(input), captured_stdout() as s:
        main()

    assert s.getvalue().strip() == 'bar'

What are provided_stdin and captured_stdout functions? They are just generators decorated by contextmanager. Here's the code:

from StringIO import StringIO
from contextlib import contextmanager
import sys


@contextmanager
def captured_stdout():
    """
    Replaces ``sys.stdout`` with StringIO object yielded to the user.
    """
    orig_stdout = sys.stdout
    buffer = StringIO()
    try:
        sys.stdout = buffer
        yield buffer
    finally:
        sys.stdout = orig_stdout


@contextmanager
def provided_stdin(io):
    """
    Replaces ``sys.stdin`` with StringIO object provided by the user.
    """
    orig_stdin = sys.stdin
    try:
        sys.stdin = io
        yield
    finally:
        sys.stdin = orig_stdin

Both functions are really super simple thanks to the abstractions provided by contextmanater.

tl;dr contextmanager is awesome, especially for writing test helper functions.