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) 

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

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