Python: Write unittest for console print

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, '1n2n')

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 == 'eggsn'

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('eggsn', captured.out)

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

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')

The answer of @Acumenus says:

It also uses a reusable helper method assert_stdout, although this helper is specific to the function being tested.

the bold part seems a big drawback, thus I would do the following instead:

# extend unittest.TestCase with new functionality
class TestCase(unittest.TestCase):

    def assertStdout(self, expected_output):
        return _AssertStdoutContext(self, expected_output)

    # as a bonus, this syntactical sugar becomes possible:
    def assertPrints(self, *expected_output):
        expected_output = "n".join(expected_output) + "n"
        return _AssertStdoutContext(self, expected_output)



class _AssertStdoutContext:

    def __init__(self, testcase, expected):
        self.testcase = testcase
        self.expected = expected
        self.captured = io.StringIO()

    def __enter__(self):
        sys.stdout = self.captured
        return self

    def __exit__(self, exc_type, exc_value, tb):
        sys.stdout = sys.__stdout__
        captured = self.captured.getvalue()
        self.testcase.assertEqual(captured, self.expected)

this allows for the much nicer and much more re-usable:

# in a specific test case, the new method(s) can be used
class TestPrint(TestCase):

    def test_print1(self):
        with self.assertStdout("testn"):
            print("test")

by using a straight forward context manager. (It might also be desirable to append "n" to expected_output since print() adds a newline by default. See next example…)

Furthermore, this very nice variant (for an arbitrary number of prints!)

    def test_print2(self):
        with self.assertPrints("test1", "test2"):
            print("test1")
            print("test2")

is possible now.

Another variant is leaning on the logging module rather than print(). This module also has a suggestion of when to use print in the documentation:

Display console output for ordinary usage of a command line script or program

PyTest has built-in support for testing logging messages.

You can also capture the standard output of a method using contextlib.redirect_stdout:

import unittest
from contextlib import redirect_stdout
from io import StringIO

class TestMyStuff(unittest.TestCase):
    # ...
    def test_stdout(self):
        with redirect_stdout(StringIO()) as sout:
            my_command_that_prints_to_stdout()
        
        # the stream replacing `stdout` is available outside the `with`
        # you may wish to strip the trailing newline
        retval = sout.getvalue().rstrip('n')

        # test the string captured from `stdout`
        self.assertEqual(retval, "whatever_retval_should_be")

Gives you a locally scoped solution. It is also possible to capture the standard error using contextlib.redirect_stderr().