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, ignore_errors=True)

    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, a source code file, 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, ignore_errors=True)

    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, ignore_errors=True)

    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, ignore_errors=True)

    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.

Testing that certain functions are called

Now we focus ourselves on this snippet of the function we are testing:

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

It's essentially testing if our destination dir for extraction exists, and creates it if it doesn't exist yet. Now there are two ways to test this. The most obvious one is to let the function run, and then see if dest_dir is created through the os.path.exists function inside our test.

However, there is another way to test this using the mock library: we test this snippet instead by testing if os.mkdir is called exactly once. I'll be using this method of testing for this tutorial.

To start off, let's create an empty test_if_dest_dir_created function:

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, ignore_errors=True)

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

    def test_if_dest_dir_created(self):
        pass

Next, we want to make sure that os.path.exists always returns False, to ensure that os.mkdir is always called. To do that, we use the mock library:

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, ignore_errors=True)

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

    @patch.object(os.path, 'exists')
    def test_if_dest_dir_created(self, exists_method):
        exists_method.return_value = False

Next, we want to check if os.mkdir is indeed called when os.path.exists is False (which we ensured in the previous line that it always will be for that test). To be able to check that, we first have to mock the os.mkdir function too. Mocks of functions provide a assert_called_once_with function that tests if the function has been called once and with the given arguments. With that, let's complete our test:

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, ignore_errors=True)

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

    @patch.object(os, 'mkdir')
    @patch.object(os.path, 'exists')
    def test_if_dest_dir_created(self, exists_method, mkdir_method):
        exists_method.return_value = False

        system_specific_unzipping('nonexistent.zip', self.dest_dir)

        mkdir_method.assert_called_once_with(self.dest_dir)

Take note of the special ordering between our @patch.object decorators and the mock method parameters. Mock method parameters follow the patch decorator sequence in reverse: whatever patch decorator we have on top becomes the last parameter in our test, and so on and so forth. When we run that, we get this output:

....E.
======================================================================
ERROR: test_if_dest_dir_created (__main__.SystemSpecificUnzippingTests)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/aron/Dropbox/src/virtualenvs/utils/local/lib/python2.7/site-packages/mock.py", line 1201, in patched
    return func(*args, **keywargs)
  File "testing/platforms_test.py", line 65, in test_if_dest_dir_created
    system_specific_unzipping('nonexistent.zip', 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 6 tests in 0.004s

FAILED (errors=1)

Damn! That exception we tested previously is interfering with this test now. Since it is not in the scope of this test to see that our function runs correctly, we can ignore the exception that is raised her to allow us to proceed to the last line, which contains our assert 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, ignore_errors=True)

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

    @patch.object(os, 'mkdir')
    @patch.object(os.path, 'exists')
    def test_if_dest_dir_created(self, exists_method, mkdir_method):
        exists_method.return_value = False

        try:
            system_specific_unzipping('nonexistent.zip', self.dest_dir)
        except:
            pass

        mkdir_method.assert_called_once_with(self.dest_dir)

Run that and we get this console output:

......
----------------------------------------------------------------------
Ran 6 tests in 0.003s

OK

Hooray! Everything works perfectly.

I think those should be enough for now. Knowledge from both Part 1 and Part 2 should allow you to be able to test a majority of python functions' behavior.

Here's a summary of things we learned so far:

• Test that exceptions are raised using self.assertRaises
• __file__ points to the current source file
• setUp and tearDown is run on every test
• We can chain multiple patch.objects for one test. They follow a LIFO sequence: the top patch.object corresponds to the last parameter, the second topmost corresponds to the second to the last parameter etc. Up until we get to the bottom patch.object, we corresponds to the second parameter of the test (the first is self).
• We test if a mocked function is called once with the assert_called_once_with method of the stubbed function.

Here's the testing/platforms_test.py we have after going through Part 1 and 2:

import os
import platform
import shutil
import sys
import unittest

import mock
from mock import patch

sys.path += [os.path.realpath('..'), os.path.realpath('.')]

from platforms import system_script_extension, system_specific_unzipping

class SystemScriptExtensionTests(unittest.TestCase):

    @patch.object(platform, 'system')
    def test_returns_bat_on_windows(self, system_method):
        system_method.return_value = 'Windows'

        self.assertEquals(system_script_extension(), '.bat')
        self.assertEquals(system_script_extension('Windows'), '.bat')

    @patch.object(platform, 'system')
    def test_returns_command_on_darwin(self, system_method):
        system_method.return_value = 'Darwin' # the Mac kernel

        self.assertEquals(system_script_extension(), '.command')
        self.assertEquals(system_script_extension('Darwin'), '.command')

    @patch.object(platform, 'system')
    def test_returns_sh_on_linux(self, system_method):
        system_method.return_value = 'Linux'

        self.assertEquals(system_script_extension(), '.sh')
        self.assertEquals(system_script_extension('Linux'), '.sh')

    @patch.object(platform, 'system')
    def test_returns_sh_by_default(self, system_method):
        system_method.return_value = 'Random OS'

        self.assertEquals(system_script_extension(), '.sh')
        self.assertEquals(system_script_extension('Random OS'), '.sh')


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, ignore_errors=True)

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

    @patch.object(os, 'mkdir')
    @patch.object(os.path, 'exists')
    def test_if_dest_dir_created(self, exists_method, mkdir_method):
        exists_method.return_value = False

        try:
            system_specific_unzipping('nonexistent.zip', self.dest_dir)
        except:
            pass

        mkdir_method.assert_called_once_with(self.dest_dir)


if __name__ == '__main__':
    unittest.main()

See you on Part 3!