David Janes' Code Weblog

When developing Python code there’s a tendency to do add a __main__ section to test the code:

def add(a, b):
    return  a + b

if __name__ == '__main__':
    print add(3, 4)

Don’t. Python has a great little package called unittest that let’s you quickly frame functions in testcases.

If the example above is called add.py, I’ll generally make a subdirectory called tests and add a test program called test_add.py. This can be as simple as:

import unittest

class TestAdd(unittest.TestCase):
    def setUp(self):
        pass

    def test_1(self):
        self.assertEqual(add(3, 4), 7)
        self.assertEqual(add(4, 4),          self.assertEqual(add(4, -4), 0)

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

But I prefer to use the following pattern:

class TestAdd(unittest.TestCase):
    def test_add(self):
        checkds = [
            {
                "a" : 4,
                "b" : 3,
                "@result": 7
            },
            {
                "a" : 4,
                "b" : 4,
                "@result": 8
            },
            {
                "a" : 4,
                "b" : -4,
                "@result": 0
            },
        ]

        for checkd in checkds:
            expected_result = checkd.pop("@result")
            actual_result = add(**checkd)

            if expected_result == -1:
                print checkd, actual_result
                continue

            try:
                self.assertEqual(expected_result, actual_result)
            except:
                print checkd, actual_result
                raise

In particular:

• the individual tests are defined in the checkds list of dictionaries
• the bottom part (the for loop) is boilerplate
  • it removes the @result from the dictionary
  • it calls add with the remaining dictionary
  • and it then asserts that the actual_result was the same as the expected_result
• if the expected_result is -1, it doesn’t run the test, it just prints the actual_result. This is great for setting up your tests in the first place. Obviously you might way to change this marker for testing functions that can return -1, but you get the idea

The advantage of using unittest like is that you’re now not depending on visual inspection or remembering which files you put a __main__ in to test your code. As a secondary benefit, unittest helps you think about edge cases, how other people might call your code.

Just go to your test directory and run them all and you’ll be sure your libraries are behaving as designed.