Beyond assert

Last time, we built a factorial function in a test-driven style using Python’s assert statement:

def factorial(n):
    if n < 0:
        raise ValueError('Factorial of negative is undefined')
    return n * factorial(n-1) if n else 1
assert factorial(0) == 1
assert factorial(2) == 2
assert factorial(5) == 120
try:
    factorial(-1)
    assert False, 'Factorial of < 0 should have raised an error'
except ValueError as e:
    assert str(e) == 'Factorial of negative is undefined'

This is not, however, the way that assert statements normally appear in code; more often, they’re woven into the program as a way to validate invariants. For example, it might not be important to us what type of exception our factorial function throws when given negative numbers, so we can simplify a little:

def factorial(n):
      assert n >= 0
      return n * factorial(n-1) if n else 1

By contrast, when writing tests, instead of writing “assert” in the module we’re testing, we more often pull the tests out of the main body of the program and run them with a testing tool. Doing this with our factorial example gives us two files, one for the code being tested and one for the tests; we’ll call them factorial.py and test_factorial.py.

The unittest Module

If we copy our first test into test_factorial.py, it looks like this:

from factorial import factorial
assert factorial(0) == 1

Rewinding to the very beginning — before we wrote any code at all in our factorial function — we would run this test and see it fail:

python test_factorial.py
Traceback (most recent call last):
  File "test_factorial.py", line 3, in <module>
    assert factorial(0) == 1
AssertionError

Plain asserts excel at testing small programs, but as our programs get more complicated, we find testing by plain asserts doesn’t give us as much information as we’d like. If we instead translate this test to run via Python’s unittest module, the test looks like this:

import unittest
from factorial import factorial
class TestFactorial(unittest.TestCase):
    def test_nothing(self):
        self.assertEqual(factorial(0), 1)
unittest.main()

Running this version tells us more about both failures and successes. Most importantly, using assertEqual() instead of the assert statement automatically tells us what values went on each side of the equality. We can now see both which assertion failed and why it failed; here, that’s because our factorial function returned None:

$ python test_factorial.py
F
====================================================================
FAIL: test_nothing (__main__.TestFactorial)
--------------------------------------------------------------------
Traceback (most recent call last):
  File "test_factorial.py", line 7, in test_nothing
    self.assertEqual(factorial(0), 1)
AssertionError: None != 1
--------------------------------------------------------------------
Ran 1 test in 0.001s
FAILED (failures=1)

If we fix our function so that this test passes, unittest tells us all is well:

$ python test_factorial.py
.
--------------------------------------------------------------------Ran 1 test in 0.000s
OK

We could translate our other tests into unittest versions in different ways. For example, we could stuff all our assertions into one function:

class TestFactorial(unittest.TestCase):
    def test_everything(self):
        self.assertEqual(factorial(0), 1)
        self.assertEqual(factorial(2), 2)
        self.assertEqual(factorial(5), 120)

Or we could split them into individual functions:

class TestFactorial(unittest.TestCase):
    def test_base_case(self):
        self.assertEqual(factorial(0), 1)
    def test_first_recursive_case(self):
        self.assertEqual(factorial(2), 2)
    def test_recursing_further(self):
        self.assertEqual(factorial(5), 120)

When we split asserts across more functions, we do so at the cost of more code. More code is bad, all else equal, so what have we gained?

Consider what happens when we rewrite our factorial function to an iterative version and suppose we work in a team that has a strange rule that all variables must start at zero, so we write it like so:

def factorial(n):
    acc = 0
    for x in range(n):
        acc = (acc or 1) * (x + 1)
    return acc

We run our tests and discover we made an error:

python test_factorial.py
F..
====================================================================
FAIL: test_base_case (__main__.TestFactorial)
--------------------------------------------------------------------
Traceback (most recent call last):
  File "test_factorial.py", line 7, in test_base_case
    self.assertEqual(factorial(0), 1)
AssertionError: 0 != 1
--------------------------------------------------------------------
Ran 3 tests in 0.001s
FAILED (failures=1)

When we were using ordinary asserts, the first failure would have stopped the test run so we wouldn’t have known whether the other cases were still ok. Here, however, unittest ran all three tests and only one failed, so it tells us that our base case is in error but our other cases are ok.

Compare with what would have happened if we had made a bigger error that caused all tests to fail:

$ python test_factorial.py
FFF
====================================================================
FAIL: test_base_case (__main__.TestFactorial)
...
====================================================================
FAIL: test_first_recursive_case (__main__.TestFactorial)
...
====================================================================
FAIL: test_recursing_further (__main__.TestFactorial)
...
Ran 3 tests in 0.001s
FAILED (failures=3)

When we were using three asserts in a row instead of in separate functions, both errors would have looked the same. When we instead split our assertions into separate functions, the test runner helps us pinpoint where we went wrong and correct our mistake faster, one of the most important ways unittest improves on plain assert statements.

We’ll examine a couple more vital test runner features next time.