Guide: Writing unit tests Part 2

Nice to see you again! I'll assume you've been through Part 1 already. If not, go there first! We're picking up where we left off from there.

Functions with exceptions

The last function we created tests for was really simple. For this guide, we're gonna test something more complex. Take a look at this function!

def system_specific_unzipping(zip_file, dest_dir, callback=_default_callback_unzip):
    """
    # unpack the inner zip to the destination
    """

    if not os.path.exists(dest_dir):
        os.mkdir(dest_dir)

    if not is_zipfile(zip_file):
        raise Exception("bad zip file")

    zip = ZipFile(zip_file, "r")
    nfiles = len(zip.namelist())

    for fi, afile in enumerate(zip.namelist()):
        if callback:
            callback(afile, fi, nfiles)

        zip.extract(afile, path=dest_dir)
        # If it's a unix script or manage.py, give permissions to execute
        if (not is_windows()) and (os.path.splitext(afile)[1] in system_specific_scripts() or afile.endswith("manage.py")):
            os.chmod(os.path.realpath(dest_dir + "/" + afile), 0775)

It's got more arguments, raises exceptions, and does some I/O. Let's get to testing this!

Testing the exceptions

Now normally, we want to test "the Happy Path" first, a.k.a. when everything goes well and the function gives a proper output. But for the sake of this tutorial, let's go ahead and test our function's exception-raising behavior. Looking at the function, this is the line we want to test:

    if not is_zipfile(zip_file):
        raise Exception("bad zip file")

Go ahead and create another test class in platform_tests.py called SystemSpecificUnzippingTests, and have a function in there called test_raises_exception_on_invalid_zipfile:

class SystemSpecificUnzippingTests(unittest.TestCase):

    def test_raises_exception_on_invalid_zipfile(self):
        pass

Reading the function again, we know that it's gonna be creating directories due to this if statement:

    if not os.path.exists(dest_dir):
        os.mkdir(dest_dir)

In order to make our directory clean before and after running each test, we add in setUp and tearDown methods which set the dest_dir and makes sure it's deleted for each test case:

class SystemSpecificUnzippingTests(unittest.TestCase):

    def setUp(self):
        self.dest_dir = os.path.join(os.path.dirname(__file__), 'extract_dir')

    def tearDown(self):
        shutil.rmtree(self.dest_dir)

    def test_raises_exception_on_invalid_zipfile(self):
        pass

Don't forget to add a import shutil line in our import statements.

Now, what we want to test here is if it raises an exception somehow. Now, what file do we know is not a zip file? Why, our current file of course, represented by __file__! Let's add in a call to system_specific_unzipping and have that as our zip file:

class SystemSpecificUnzippingTests(unittest.TestCase):

    def setUp(self):
        self.dest_dir = os.path.join(os.path.dirname(__file__), 'extract_dir')

    def tearDown(self):
        shutil.rmtree(self.dest_dir)

    def test_raises_exception_on_invalid_zipfile(self):
        system_specific_unzipping(__file__, self.dest_dir)

When we run our tests with python testing/platforms_test.py, we get:

....E
======================================================================
ERROR: test_raises_exception_on_invalid_zipfile (__main__.SystemSpecificUnzippingTests)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "testing/platforms_test.py", line 57, in test_raises_exception_on_invalid_zipfile
    system_specific_unzipping(__file__, self.dest_dir)
  File "/home/aron/Dropbox/src/fle-utils/platforms.py", line 127, in system_specific_unzipping
    raise Exception("bad zip file")
Exception: bad zip file

----------------------------------------------------------------------
Ran 5 tests in 0.002s

FAILED (errors=1)

Hooray! An exception is raised, which was what we wanted to test! However, we have to tell python that us getting an exception is a good thing. To do that, we wrap our system_specific_unzipping inside an assertRaises call:

class SystemSpecificUnzippingTests(unittest.TestCase):

    def setUp(self):
        self.dest_dir = os.path.join(os.path.dirname(__file__), 'extract_dir')

    def tearDown(self):
        shutil.rmtree(self.dest_dir)

    def test_raises_exception_on_invalid_zipfile(self):
        self.assertRaises(Exception, system_specific_unzipping, __file__, self.dest_dir)

Take note that the way we call system_specific_unzipping now changes. assertRaises can also be used as a context manager too, inside a with statement. This test can be rephrased as:

class SystemSpecificUnzippingTests(unittest.TestCase):

    def setUp(self):
        self.dest_dir = os.path.join(os.path.dirname(__file__), 'extract_dir')

    def tearDown(self):
        shutil.rmtree(self.dest_dir)

    def test_raises_exception_on_invalid_zipfile(self):
        with self.assertRaises(Exception):
            system_specific_unzipping(__file__, self.dest_dir)

I personally prefer using it as a with statement.

Whatever way you write the test, you should get this output:

.....
----------------------------------------------------------------------
Ran 5 tests in 0.002s

OK

Awesome! Python now knows that an exception being raised here is a good thing, and lets the test pass.