Unit Testing with unittest

A unit test checks a small piece of behavior against an expected result. Python’s standard-library module unittest groups related test methods in classes and reports which checks pass or fail. We will build complete test files, then add setup and cleanup and combine tests into a suite.

Assertions and Test Cases

An assertion compares an actual result with an expectation. Subclass unittest.TestCase and write methods beginning with test_; each method becomes a separate test case. Useful checks include assertEqual, assertTrue, assertFalse, assertIsNone, assertIsInstance, assertIn and assertIs.

Save the following as math_helpers.py in a new practice folder:

# math_helpers.py
def add(a, b):
    return a + b

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

The module supplies functions; defining them produces no output. In the same folder, save test_math_helpers.py:

# test_math_helpers.py
import unittest
from math_helpers import add, safe_divide

class TestMathHelpers(unittest.TestCase):
    def test_add(self):
        self.assertEqual(add(2, 3), 5)
        self.assertNotEqual(add(2, 2), 5)

    def test_safe_divide(self):
        self.assertEqual(safe_divide(10, 2), 5)
        self.assertRaises(ValueError, safe_divide, 10, 0)

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

assertRaises calls safe_divide(10, 0) and checks that it raises ValueError. The if __name__ == "__main__" guard starts the runner when this file is executed directly, not when another file imports its classes.

Run the Tests

From the folder containing both files, run:

python -m unittest test_math_helpers.py

Output — elapsed time varies:

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

OK

The dots represent successful tests. OK means the assertions in those two tests passed; it does not establish correctness for untested inputs. Add -v for the names of individual tests.

Fixtures: Setup and Cleanup

A fixture prepares the data or resources a test needs. setUp() runs before a test method; tearDown() runs afterward if setup succeeded. addCleanup() registers cleanup that also runs if a later step in setup fails.

Save this function as file_tools.py:

# file_tools.py
def read_first_line(path):
    with open(path, "r", encoding="utf-8") as file:
        return file.readline().strip()

Now save test_file_tools.py. A temporary directory keeps the test from overwriting a file in your project:

# test_file_tools.py
import tempfile
import unittest
from pathlib import Path
from file_tools import read_first_line

class TestFileTools(unittest.TestCase):
    def setUp(self):
        self.directory = tempfile.TemporaryDirectory()
        self.addCleanup(self.directory.cleanup)
        self.test_file = Path(self.directory.name) / "data.txt"
        self.test_file.write_text("Hello\nSecond line\n", encoding="utf-8")

    def test_read_first_line(self):
        self.assertEqual(read_first_line(self.test_file), "Hello")

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

TemporaryDirectory() creates a separate folder. Cleanup removes it after the test, while read_first_line closes its file through with. Run the test:

python -m unittest test_file_tools.py

Output — elapsed time varies:

.
----------------------------------------------------------------------
Ran 1 test in ...s

OK

More Test Cases

A test class can group several related behaviors. Save test_strings.py:

# test_strings.py
import unittest

class TestStringMethods(unittest.TestCase):
    def test_upper(self):
        self.assertEqual("Ada".upper(), "ADA")

    def test_lower(self):
        self.assertEqual("ADA".lower(), "ada")

    def test_split(self):
        self.assertEqual("Ada Grace".split(), ["Ada", "Grace"])

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

Each method checks one operation. Run python -m unittest test_strings.py to execute its three tests, or use discovery after saving all the example files:

python -m unittest discover -v

Output summary — individual test names and elapsed time vary with the selected files:

Ran 6 tests in ...s

OK

By default, discovery searches for importable files matching test*.py. Here it finds two math tests, one file test and three string tests. It does not automatically select every file ending in _test.py.

Test Suites and Runners

A suite groups selected tests. A runner executes that suite and reports results. Save run_suite.py alongside the files above:

# run_suite.py
import unittest
from test_math_helpers import TestMathHelpers
from test_strings import TestStringMethods

loader = unittest.defaultTestLoader
suite = unittest.TestSuite([
    loader.loadTestsFromTestCase(TestMathHelpers),
    loader.loadTestsFromTestCase(TestStringMethods),
])

runner = unittest.TextTestRunner()
runner.run(suite)

Run it with python run_suite.py. This suite deliberately selects the two math tests and three string tests:

Output — elapsed time varies:

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

OK

For a notebook, save these files first and run python -m unittest through that notebook’s interpreter. If test classes are defined in the notebook itself, call unittest.main(argv=[""], exit=False) after their definitions; the adjusted arguments avoid the notebook’s startup flags and exit=False keeps its kernel running.

Try it: change one expected string to an incorrect value, run that test and read the failure. Restore it, then add a test for a missing file using assertRaises(FileNotFoundError, ...).

The same behavior can also be checked with ordinary test functions. The next section uses pytest to do that.