Python: Write unittest for console print

You can easily capture standard output by just temporarily redirecting sys.stdout to a StringIO object, as follows:

import StringIOimport sysdef foo(inStr):    print "hi"+inStrdef test_foo():    capturedOutput = StringIO.StringIO()          # Create StringIO object    sys.stdout = capturedOutput                   #  and redirect stdout.    foo('test')                                   # Call unchanged function.    sys.stdout = sys.__stdout__                   # Reset redirect.    print 'Captured', capturedOutput.getvalue()   # Now works as before.test_foo()

The output of this program is:

Captured hitest

showing that the redirection successfully captured the output and that you were able to restore the output stream to what it was before you began the capture.

Note that the code above in for Python 2.7, as the question indicates. Python 3 is slightly different:

import ioimport sysdef foo(inStr):    print ("hi"+inStr)def test_foo():    capturedOutput = io.StringIO()                  # Create StringIO object    sys.stdout = capturedOutput                     #  and redirect stdout.    foo('test')                                     # Call function.    sys.stdout = sys.__stdout__                     # Reset redirect.    print ('Captured', capturedOutput.getvalue())   # Now works as before.test_foo()

This Python 3 answer uses unittest.mock. It also uses a reusable helper method assert_stdout, although this helper is specific to the function being tested.

import ioimport unittestimport unittest.mockfrom .solution import fizzbuzzclass TestFizzBuzz(unittest.TestCase):    @unittest.mock.patch('sys.stdout', new_callable=io.StringIO)    def assert_stdout(self, n, expected_output, mock_stdout):        fizzbuzz(n)        self.assertEqual(mock_stdout.getvalue(), expected_output)    def test_only_numbers(self):        self.assert_stdout(2, '1\n2\n')

Note that the mock_stdout arg is passed automatically by the unittest.mock.patch decorator to the assert_stdout method.

A general-purpose TestStdout class, possibly a mixin, can in principle be derived from the above.

For those using Python ≥3.4, contextlib.redirect_stdout also exists, but it seems to serve no benefit over unittest.mock.patch.

If you happen to use pytest, it has builtin output capturing. Example (pytest-style tests):

def eggs():    print('eggs')def test_spam(capsys):    eggs()    captured = capsys.readouterr()    assert captured.out == 'eggs\n'

You can also use it with unittest test classes, although you need to passthrough the fixture object into the test class, for example via an autouse fixture:

import unittestimport pytestclass TestSpam(unittest.TestCase):    @pytest.fixture(autouse=True)    def _pass_fixtures(self, capsys):        self.capsys = capsys    def test_eggs(self):        eggs()        captured = self.capsys.readouterr()        self.assertEqual('eggs\n', captured.out)

Check out Accessing captured output from a test function for more info.