Testing with assert and unittest

Objectives

By the end of this chapter, you should be able to:

• Explain why automated tests matter
• Use assert to check a condition
• Write a basic test case with the unittest module

💡 Why this matters: Manually re-running a script and eyeballing the output doesn’t scale, and it’s easy to miss a case you broke without noticing. A test suite checks your assumptions automatically, every time, so regressions get caught immediately instead of in production.

assert

The simplest form of testing in Python is the assert statement. It does nothing if the condition is truthy, and raises an AssertionError if it isn’t:

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

assert add(2, 2) == 4          # passes silently
assert add(2, 2) == 5, "2 + 2 should be 4"  # AssertionError: 2 + 2 should be 4

The optional second argument becomes the error message, which makes a failing assertion far easier to diagnose than a bare one. assert is useful for quick sanity checks, but it’s not a full testing framework: it doesn’t organize test cases, run a whole suite at once, or report which of many checks failed.

unittest

Python’s built-in unittest module gives you that structure. You define a class inheriting from unittest.TestCase, and each method starting with test_ is a separate, independently-run test:

import unittest

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

class TestAdd(unittest.TestCase):
    def test_adds_two_positive_numbers(self):
        self.assertEqual(add(2, 2), 4)

    def test_adds_a_negative_number(self):
        self.assertEqual(add(5, -3), 2)

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

Running this file directly (python3 test_add.py) executes both test_ methods and reports a pass/fail summary. If one fails, you’ll see exactly which one and why, without the others being affected:

..
----------------------------------------------------------------------
Ran 2 tests in 0.000s

OK

Common Assertion Methods

unittest.TestCase gives you far more than just equality checks:

That last one is worth a closer look: it’s how you test that your code fails correctly:

def divide(a, b):
    if b == 0:
        raise ValueError("Cannot divide by zero")
    return a / b

class TestDivide(unittest.TestCase):
    def test_raises_on_zero_division(self):
        self.assertRaises(ValueError, divide, 10, 0)

Try It

1. Write a function and a handful of assert statements checking it against a few inputs, including at least one edge case.
2. Convert those assertions into a unittest.TestCase with individually named test_ methods.
3. Write a function that raises an error under some condition, then use assertRaises to confirm it does.

Recap

• assert is a quick, single-condition check that raises AssertionError (with an optional message) when it fails.
• unittest.TestCase organizes multiple related tests into a class, running each test_ method independently and reporting which ones fail.
• assertEqual, assertTrue, assertIn, and assertRaises cover most everyday testing needs.

Next lesson: reading command-line arguments with sys.argv and argparse.