Python: Write unittest for console print

Question

Function foo prints to console. I want to test the console print. How can I achieve this in python?

Need to test this function, has NO return statement :

def foo(inStr):
   print "hi"+inStr

My test :

def test_foo():
    cmdProcess = subprocess.Popen(foo("test"), stdout=subprocess.PIPE)
    cmdOut = cmdProcess.communicate()[0]
    self.assertEquals("hitest", cmdOut)

Solution 1

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

import StringIO
import sys

def foo(inStr):
    print "hi"+inStr

def 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 io
import sys

def 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()

Solution 2

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 io
import unittest
import unittest.mock

from .solution import fizzbuzz


class 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.

Solution 3

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 unittest
import pytest


class 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.

Solution 4

You can also use the mock package as shown below, which is an example from https://realpython.com/lessons/mocking-print-unit-tests.

from mock import patch

def greet(name):
    print('Hello ', name)

@patch('builtins.print')
def test_greet(mock_print):
    # The actual test
    greet('John')
    mock_print.assert_called_with('Hello ', 'John')
    greet('Eric')
    mock_print.assert_called_with('Hello ', 'Eric')

View more solutions