Unittesting print statements

Recently I was refactoring a small package that is supposed to allow execution of arbitrary python code on a remote machine. The first implementation was working nicely but with one serious drawback – function handling the actual code execution was running in a synchronous (blocking) mode. As the result all of the output (both stdout and stderr) was presented only at the end, i.e. when code finished its execution. This was unacceptable since the package should work in a way as transparent to the user as possible. So a wall of text when code completes its task wasn’t acceptable.

The goal of the refactoring was simple – to have the output presented to the user immediately after it was printed on the remote host. As a TDD worshipper I wanted to start this in a kosher way, i.e. with a test. And I got stuck.

For a day or so I had no idea how to progress. How do you unittest the print statements? It’s funny when I think about this now. I have used a similar technique many times in the past for output redirection, yet somehow haven’t managed to make a connection with this problem.

The print statement

So how do you do it? First we should understand what happens when print statement is executed. In python 2.x the print statement does two things – converts provided expressions into strings and writes the result to a file like object handling the stdout. Conveniently it is available as sys.stdout (i.e. as a part of sys module). So all you have to do is to overwrite the sys.stdout with your own object providing a ‘write’ method. Later you may discover, that some other methods may be also needed (e.g. ‘flush’ is quite often used), but for starters, having only the ‘write’ method should be sufficient.

A first try – simple stdout interceptor

The code below does just that. The MyOutput class is designed to replace the original sys.stdout:

import unittest
import sys

def fn_print(nrepeat):
    print "ab"*nrepeat

class MyTest(unittest.TestCase):
    def test_stdout(self):
        class MyOutput(object):
            def __init__(self):
                self.data = []

            def write(self, s):
                self.data.append(s)

            def __str__(self):
                return "".join(self.data)

        stdout_org = sys.stdout
        my_stdout = MyOutput()
        try:
            sys.stdout = my_stdout
            fn_print(2)
        finally:
            sys.stdout = stdout_org

        self.assertEquals( str(my_stdout), "abab\n") 

if __name__ == "__main__":
    unittest.main()

The fn_print function provides output to test against. After replacing sys.stdout we call this function and compare the obtained output with the expected one. It is worth noting that in the example above the original sys.stdout is first preserved and then carefully restored inside the ‘finally’ block. If you don’t do this you are likely to loose any output coming from other tests.

Is my code async? Logging time of arrival

In the second example we will address the original problem – is output presented as a wall of text at the end or maybe in real time as we want to. For this we will add time of arrival logging capability to the object replacing sys.stdout:

import unittest
import time
import sys

def fn_print_with_delay(nrepeat):
    for i in xrange(nrepeat):
        print    # prints a single newline
        time.sleep(0.5)

class TestServer(unittest.TestCase):
    def test_stdout_time(self):
        class TimeLoggingOutput(object):
            def __init__(self):
                self.data = []
                self.timestamps = []

            def write(self, s):
                self.timestamps.append(time.time())
                self.data.append(s)

        stdout_org = sys.stdout
        my_stdout = TimeLoggingOutput()
        nrep = 3 # make sure is >1
        try:
            sys.stdout = my_stdout
            fn_print_with_delay(nrep)
        finally:
            sys.stdout = stdout_org

        for i in xrange(nrep):
            if i > 0:
                dt = my_stdout.timestamps[i]-my_stdout.timestamps[i-1]
                self.assertTrue(0.5<dt<0.52)

if __name__ == "__main__":
    unittest.main()

The code is pretty much self-explanatory – the fn_print_with_delay function prints newlines in half of a second intervals. We override sys.stdout with an instance of a class capable of storing timestamps (obtained with time.time()) of all calls to the write method. At the and we assert the timestamps are spaced half of a second approximately. The code above works as expected:

.
----------------------------------------------------------------------
Ran 1 test in 1.502s

OK

If we change the interval inside the fn_print_with_delay function to one second, the test will (fortunately) fail.

Wrap-up

As we saw, testing for expected output is in fact trivial – all you have to do is to put an instance of a class with a ‘write’ method in proper place (i.e. sys.stdout). The only ‘gotcha’ is the cleanup – you should remember to restore sys.stdout to its original state. You may apply the exact same technique if you need to test stderr (just target the sys.stderr instead of sys.stdout). It is also worth noting that using a similar technique you could intercept (or completely silence) output coming from external libraries.