How to do unit testing of functions writing files using Python’s ‘unittest’

I prefer to have output functions explicitly accept a file handle (or file-like object), rather than accept a file name and opening the file themselves. This way, I can pass a StringIO object to the output function in my unit test, then .read() the contents back from that StringIO object (after a .seek(0) call) and compare with my expected output.

For example, we would transition code like this

##File:lamb.py
import sys


def write_lamb(outfile_path):
    with open(outfile_path, 'w') as outfile:
        outfile.write("Mary had a little lamb.\n")


if __name__ == '__main__':
    write_lamb(sys.argv[1])



##File test_lamb.py
import unittest
import tempfile

import lamb


class LambTests(unittest.TestCase):
    def test_lamb_output(self):
        outfile_path = tempfile.mkstemp()[1]
        try:
            lamb.write_lamb(outfile_path)
            contents = open(tempfile_path).read()
        finally:
            # NOTE: To retain the tempfile if the test fails, remove
            # the try-finally clauses
            os.remove(outfile_path)
        self.assertEqual(contents, "Mary had a little lamb.\n")

to code like this

##File:lamb.py
import sys


def write_lamb(outfile):
    outfile.write("Mary had a little lamb.\n")


if __name__ == '__main__':
    with open(sys.argv[1], 'w') as outfile:
        write_lamb(outfile)



##File test_lamb.py
import unittest
from io import StringIO

import lamb


class LambTests(unittest.TestCase):
    def test_lamb_output(self):
        outfile = StringIO()
        # NOTE: Alternatively, for Python 2.6+, you can use
        # tempfile.SpooledTemporaryFile, e.g.,
        #outfile = tempfile.SpooledTemporaryFile(10 ** 9)
        lamb.write_lamb(outfile)
        outfile.seek(0)
        content = outfile.read()
        self.assertEqual(content, "Mary had a little lamb.\n")

This approach has the added benefit of making your output function more flexible if, for instance, you decide you don’t want to write to a file, but some other buffer, since it will accept all file-like objects.

Note that using StringIO assumes the contents of the test output can fit into main memory. For very large output, you can use a temporary file approach (e.g., tempfile.SpooledTemporaryFile).