Testing Python Scripts

Question

How do I test the STDOUT output of a Python script with a testing framework like doctest, unittest, nose, etc? For example, say running my script "todo.py --list" should return "take out the garbage". I've read someone who separates out the STDOUT printing part of the script from the part that generates the output to be printed. I'm used to sprinkling print statements all around my shell scripts. Is this simply a TDD unfriendly habit I should break or is there a way to easily test for correct print output?

Answer 1

I see two ways :

1. Redirect stdout during the unittest:

   class YourTest(TestCase):
       def setUp(self):
           self.output = StringIO()
           self.saved_stdout = sys.stdout
           sys.stdout = self.output
   
       def tearDown(self):
           self.output.close()
           sys.stdout = self.saved_stdout
   
       def testYourScript(self):
           yourscriptmodule.main()
           assert self.output.getvalue() == "My expected ouput"
   

2. Use a logger for your outputs and listen to it in your test.

Answer 2

Python's own test suite does this quite a bit, and we use two main techniques:

1. Redirecting stdout (as others have suggested). We use a context manager for this:

   import io
   import sys
   import contextlib
   
   @contextlib.contextmanager
   def captured_output(stream_name):
       """Run the 'with' statement body using a StringIO object in place of a
          specific attribute on the sys module.
          Example use (with 'stream_name=stdout'):
   
          with captured_stdout() as s:
              print("hello")
              assert s.getvalue() == "hello"
       """
       orig_stdout = getattr(sys, stream_name)
       setattr(sys, stream_name, io.StringIO())
       try:
           yield getattr(sys, stream_name)
       finally:
           setattr(sys, stream_name, orig_stdout)
   
   def captured_stdout():
       return captured_output("stdout")
   
   def captured_stderr():
       return captured_output("stderr")
   
   def captured_stdin():
       return captured_output("stdin")
   

2. Using the subprocess module. We use this when we specifically want to test handling of command line arguments. See http://hg.python.org/cpython/file/default/Lib/test/test_cmd_line_script.py for several examples.

Answer 3

when you use py.test for your testing. You can use the "capsys" or the "capfd" test function arguments to run asserts against STDOUT and STDIN

def test_myoutput(capsys): # or use "capfd" for fd-level
    print ("hello")
    sys.stderr.write("world\n")
    out, err = capsys.readouterr()
    assert out == "hello\n"
    assert err == "world\n"
    print "next"
    out, err = capsys.readouterr()
    assert out == "next\n"

More details can be found in the py.test docs

Answer 4

Here is something that I wrote one evening that tests script runs. Note that the test does cover the basic cases, but it is not thorough enough to be a unittest by itself. Consider it a first draft.

import sys
import subprocess

if sys.platform == "win32":
   cmd = "zs.py"
else:
   cmd = "./zs.py"

def testrun(cmdline):
   try:
      retcode = subprocess.call(cmdline, shell=True)
      if retcode < 0:
         print >>sys.stderr, "Child was terminated by signal", -retcode
      else:
         return retcode
   except OSError, e:
      return e

tests = []
tests.append( (0, " string pattern 4") )
tests.append( (1, " string pattern") )
tests.append( (3, " string pattern notanumber") )
passed = 0

for t in tests:
   r = testrun(cmd + t[1])
   if r == t[0]:
      res = "passed"
      passed += 1
   else:
      res = "FAILED"
   print res, r, t[1]

print
if passed != len(tests):
   print "only",passed,"tests passed"
else:
   print "all tests passed"

And here is the script that was being tested, zs.py, This does pattern searches in a string similar to the way biochemists search for patterns in DNA data or protein chain data.

#!/usr/bin/env python

# zs - some example Python code to demonstrate to Z??s
#      interviewers that the writer really does know Python

import sys
from itertools import *

usage = '''
   Usage: zs <string> <pattern> <n>"
          print top n matches of pattern in substring"
'''

if sys.hexversion > 0x03000000:
   print "This script is only intended to run on Python version 2"
   sys.exit(2)

if len(sys.argv) != 4:
   print usage
   sys.exit(1)

A = sys.argv[1] # string to be searched
B = sys.argv[2] # pattern being searched for
N = sys.argv[3] # number of matches to report

if not N.isdigit():
   print "<n> must be a number"
   print usage
   sys.exit(3)

def matchscore(s1, s2):
   ''' a helper function to calculate the match score
   '''
   matches = 0
   for i in xrange(len(s1)):
      if s1[i] == s2[i]:
         matches += 1
   return (matches + 0.0) / len(s1)  # added 0.0 to force floating point div

def slices(s, n):
   ''' this is a generator that returns the sequence of slices of
       the input string s that are n characters long '''
   slen = len(s)
   for i in xrange(slen - n + 1):
      yield s[i:i+n]

matchlen = len(B)
allscores = ((matchscore(x,B),x,i) for i,x in enumerate(slices(A,matchlen)))
nonzeros = [ y for y in allscores if y[0] != 0 ]

for elem in sorted(nonzeros,key=lambda e: e[0],reverse=True):
   nprinted = 0 # We will count them; in case num elements > N
   print elem[1], str(round(elem[0],4)), elem[2]
   nprinted += 1
   if nprinted >= N:
      break

Answer 5

I also might want to look at the TextTest testing framework. It focusses more on functional/acceptance testing (so is less amenable to unit testing) and relies heavily on a program's textual output. This way your habit becomes a good one :-).