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()         # Make StringIO.
    sys.stdout = capturedOutput                  # Redirect stdout.
    foo('test')                                  # Call function.
    sys.stdout = sys.__stdout__                  # Reset redirect.
    print 'Captured', capturedOutput.getvalue()  # Now works.
    
test_foo()

The output of this program is:

Captured hitest

This shows 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.
    sys.stdout = capturedOutput                    # Redirect stdout.
    foo('test')                                    # Call function.
    sys.stdout = sys.__stdout__                    # Reset redirect.
    print ('Captured', capturedOutput.getvalue())  # Now works.
    
    test_foo()

---

You also need to be careful if your function can raise an exception since that may mean that the redirection reset may not be done. This may then affect later tests.

One way to solve that would be to wrap the function call in a try/except block and, on exception, reset the redirection and re-throw the exception. The single-line call to your function would then become something like this, with the other code remaining as-is:

try:
    foo('test')                  # Call function.
except:
    sys.stdout = sys.__stdout__  # Reset redirect.
    raise