Python's unittest: Writing Unit Tests for Your Code

Creating Test Cases

Before you write tests with unittest, you need some code to test. Suppose that you need to get a person’s age, process that information, and display their current life stage. For example, if the person’s age is:

• Between 0 and 9, both included, the function should return "Child".
• Greater than 9 and less than or equal to 18, the function should return "Adolescent".
• Greater than 18 and less than or equal to 65, the function should return "Adult".
• Greater than 65 and less than or equal to 150, the function should return "Golden age".
• Negative or greater than 150, the function should return "Invalid age".

In this situation, you can write a function like the following:

def categorize_by_age(age):
    if 0 <= age <= 9:
        return "Child"
    elif 9 < age <= 18:
        return "Adolescent"
    elif 18 < age <= 65:
        return "Adult"
    elif 65 < age <= 150:
        return "Golden age"
    else:
        return f"Invalid age: {age}"

This function should return correct results with different age values. To make sure that the function works correctly, you can write some unittest tests.

Following the pattern from the previous section, you’ll start by subclassing TestCase and add some methods that will help you test different input values and the corresponding results:

import unittest

from age import categorize_by_age

class TestCategorizeByAge(unittest.TestCase):
    def test_child(self):
        self.assertEqual(categorize_by_age(5), "Child")

    def test_adolescent(self):
        self.assertEqual(categorize_by_age(15), "Adolescent")

    def test_adult(self):
        self.assertEqual(categorize_by_age(30), "Adult")

    def test_golden_age(self):
        self.assertEqual(categorize_by_age(70), "Golden age")

    def test_negative_age(self):
        self.assertEqual(categorize_by_age(-1), "Invalid age: -1")

    def test_too_old(self):
        self.assertEqual(categorize_by_age(151), "Invalid age: 151")

In this example, you create a subclass of TestCase with the descriptive name TestCategorizeByAge. Note that the class name starts with Test, which is a widely used convention to make the purpose of the class immediately clear to anyone reading your code.

Also, note that the containing file is called test_age.py. By default, unittest supports test discovery based on the name of test modules. The default naming pattern is test*.py. Here, the asterisk (*) represents any sequence of characters, so starting your modules with test is recommended if you want to take advantage of the default test discovery configuration.

Then, you define six methods. Each method tests for an input value and the expected result. The methods use the .assertEqual() method from the parent class to check whether the function’s output equals the expected value.

Note that the tests above check every possible branch in the categorize_by_age() function. However, they don’t cover the boundary cases where the input age is the lower or upper limit of the interval. To make sure that the function responds as expected in those cases, you can add the following tests:

# ...

class TestCategorizeByAge(unittest.TestCase):
    # ...

    def test_boundary_child_adolescent(self):
        self.assertEqual(categorize_by_age(9), "Child")
        self.assertEqual(categorize_by_age(10), "Adolescent")

    def test_boundary_adolescent_adult(self):
        self.assertEqual(categorize_by_age(18), "Adolescent")
        self.assertEqual(categorize_by_age(19), "Adult")

    def test_boundary_adult_golden_age(self):
        self.assertEqual(categorize_by_age(65), "Adult")
        self.assertEqual(categorize_by_age(66), "Golden age")

These test methods have two assertions each. The first assertion checks for the upper limit of the age interval, and the second assertion checks for the lower limit of the next age interval.

Using multiple assertions in a test method helps you reduce boilerplate code. For example, if you use a single assertion to write these tests, then you’ll have to write six test methods instead of just three. Each method will need a unique name, which can be a challenge.

In general, using multiple assertions in your test methods has the following pros:

• Efficiency: Multiple assertions in a single test can reduce repetitive code. It can also make tests run faster in those scenarios where you have setup and teardown requirements for each test.
• Contextual testing: Multiple assertions might be necessary to check that a function behaves correctly in a specific context.
• Convenience: Multiple assertions in a test can be more straightforward and less tedious to write compared to writing multiple single-assertion tests.

The approach also has its cons:

• Clarity and isolation: When a test with multiple assertions fails, it can be harder to immediately identify which assertion caused the failure. This can go against your debugging process.
• Breakage risk: When an early assertion in a test fails, subsequent assertions are not executed. This can hide additional issues.
• Test purpose blurring: When a test has multiple assertions, it can become less focused. This can make the test harder to understand.

With the test cases in place, you’re ready to run them and see whether your categorize_by_age() function works as expected.

Running unittest Tests

Once you’ve written the tests, you need a way to run them. You’ll have at least two standard ways to run tests with unittest:

1. Make the test module executable
2. Use the command-line interface of unittest

To make a test module executable in unittest, you can add the following code to the end of the module:

# ...

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

The main() function from unittest allows you to load and run a set of tests. You can also use this function to make the test module executable. Once you’ve added these lines of code, you can run the module as a regular Python script:

$ python test_age.py
.........
----------------------------------------------------------------------
Ran 9 tests in 0.000s

OK

This command runs the tests from test_age.py. In the output, every dot represents a passing test. Then, you have a quick summary of the number of run tests and the execution time. All the tests passed, so you get an OK at the end of the output.

Among other arguments, the main() function takes the verbosity one. With this argument, you can tweak the output’s verbosity, which has three possible values:

• 0 for quiet
• 1 for normal
• 2 for detailed

Go ahead and update the call to main() as in the following snippet:

# ...

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

In the highlighted line, you set the verbosity level to 2. This update makes unittest generate a more detailed output when you run the test module:

$ python test_age.py
test_adolescent (__main__.TestCategorizeByAge.test_adolescent) ... ok
test_adult (__main__.TestCategorizeByAge.test_adult) ... ok
test_boundary_adolescent_adult (__main__.TestCategorizeByAge.test_boundary_adolescent_adult) ... ok
test_boundary_adult_golden_age (__main__.TestCategorizeByAge.test_boundary_adult_golden_age) ... ok
test_boundary_child_adolescent (__main__.TestCategorizeByAge.test_boundary_child_adolescent) ... ok
test_child (__main__.TestCategorizeByAge.test_child) ... ok
test_golden_age (__main__.TestCategorizeByAge.test_golden_age) ... ok
test_negative_age (__main__.TestCategorizeByAge.test_negative_age) ... ok
test_too_old (__main__.TestCategorizeByAge.test_too_old) ... ok

----------------------------------------------------------------------
Ran 9 tests in 0.000s

OK

This output is more detailed. It shows the tests and their result. At the end, it summarizes the test run as usual.

If you want to make the detailed output more descriptive, then you can add docstrings to your tests like in the following code snippet:

# ...

class TestCategorizeByAge(unittest.TestCase):
    def test_child(self):
        """Test for 'Child'"""
        self.assertEqual(categorize_by_age(5), "Child")

    def test_adolescent(self):
        """Test for 'Adolescent'"""
        self.assertEqual(categorize_by_age(15), "Adolescent")

    def test_adult(self):
        """Test for 'Adult'"""
        self.assertEqual(categorize_by_age(30), "Adult")

    def test_golden_age(self):
        """Test for 'Golden age'"""
        self.assertEqual(categorize_by_age(70), "Golden age")

    def test_negative_age(self):
        """Test for negative age"""
        self.assertEqual(categorize_by_age(-1), "Invalid age: -1")

    def test_too_old(self):
        """Test for too old"""
        self.assertEqual(categorize_by_age(151), "Invalid age: 151")

    def test_boundary_child_adolescent(self):
        """Test for boundary between 'Child' and 'Adolescent'"""
        self.assertEqual(categorize_by_age(9), "Child")
        self.assertEqual(categorize_by_age(10), "Adolescent")

    def test_boundary_adolescent_adult(self):
        """Test for boundary between 'Adolescent' and 'Adult'"""
        self.assertEqual(categorize_by_age(18), "Adolescent")
        self.assertEqual(categorize_by_age(19), "Adult")

    def test_boundary_adult_golden_age(self):
        """Test for boundary between 'Adult' and 'Golden age'"""
        self.assertEqual(categorize_by_age(65), "Adult")
        self.assertEqual(categorize_by_age(66), "Golden age")

# ...

In this update of your test methods, you add human-readable docstrings. Now, when you run the test with a verbosity of 2, you get the following output:

$ python test_age.py
test_adolescent (__main__.TestCategorizeByAge.test_adolescent)
Test for 'Adolescent' ... ok
test_adult (__main__.TestCategorizeByAge.test_adult)
Test for 'Adult' ... ok
test_boundary_adolescent_adult (__main__.TestCategorizeByAge.test_boundary_adolescent_adult)
Test for boundary between 'Adolescent' and 'Adult' ... ok
test_boundary_adult_golden_age (__main__.TestCategorizeByAge.test_boundary_adult_golden_age)
Test for boundary between 'Adult' and 'Golden age' ... ok
test_boundary_child_adolescent (__main__.TestCategorizeByAge.test_boundary_child_adolescent)
Test for boundary between 'Child' and 'Adolescent' ... ok
test_child (__main__.TestCategorizeByAge.test_child)
Test for 'Child' ... ok
test_golden_age (__main__.TestCategorizeByAge.test_golden_age)
Test for 'Golden age' ... ok
test_negative_age (__main__.TestCategorizeByAge.test_negative_age)
Test for negative age ... ok
test_too_old (__main__.TestCategorizeByAge.test_too_old)
Test for too old ... ok

----------------------------------------------------------------------
Ran 9 tests in 0.000s

OK

Now, apart from the odd function names, unittest uses the docstrings to improve the output’s readability, which is great.

Skipping Tests

The unittest framework also supports skipping individual test methods and even whole test case classes. Skipping tests allows you to temporarily bypass a test case without permanently removing it from your test suite.

Here are some common situations where you may need to skip tests:

• Incomplete feature: When you have an incomplete feature, you’ll need to skip the related tests to avoid false negatives.
• External services: When some of your tests depend on external services or resources that aren’t available, you may need to skip them until the service is back.
• Conditional execution: When you have tests that depend on specific conditions, such as a platform or a Python version, then you can conditionally skip the tests based on dynamic conditions.
• Known failures: When your code has a known bug, you might skip the failing test with a reference to the bug report.
• Performance considerations: When you have tests that are time-consuming or resource-intensive, you might want to skip them during regular development cycles to speed up the testing process.
• Deprecated features: When you have deprecated features that haven’t been removed yet, you can skip their related tests until the feature is removed entirely.

The following decorators will help you with the goal of skipping tests during your test running process:

In these decorators, the reason argument should describe why the test will be skipped. Consider the following toy example that shows how the decorators work:

import sys
import unittest

class SkipTestExample(unittest.TestCase):
    @unittest.skip("Unconditionally skipped test")
    def test_unimportant(self):
        self.fail("The test should be skipped")

    @unittest.skipIf(sys.version_info < (3, 12), "Requires Python >= 3.12")
    def test_using_calendar_constants(self):
        import calendar

        self.assertEqual(calendar.Month(10), calendar.OCTOBER)

    @unittest.skipUnless(sys.platform.startswith("win"), "Requires Windows")
    def test_windows_support(self):
        from ctypes import WinDLL, windll

        self.assertIsInstance(windll.kernel32, WinDLL)

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

In this sample test module, you have three test methods. The first one never runs because you use the @skip decorator on it. The second test method only runs if you’re on a Python version equal to or greater than 3.12. Finally, the last test method runs if you’re on a Windows box.

• Windows
• Linux + macOS

PS> python skip_tests.py
test_unimportant (__main__.SkipTestExample.test_unimportant) ... skipped 'Unconditionally skipped test'
test_using_calendar_constants (__main__.SkipTestExample.test_using_calendar_constants) ... skipped 'Requires Python >= 3.12'
test_windows_support (__main__.SkipTestExample.test_windows_support) ... ok

----------------------------------------------------------------------
Ran 3 tests in 0.156s

OK (skipped=2)

$ python skip_tests.py
test_unimportant (__main__.SkipTestExample.test_unimportant) ... skipped 'Unconditionally skipped test'
test_using_calendar_constants (__main__.SkipTestExample.test_using_calendar_constants) ... ok
test_windows_support (__main__.SkipTestExample.test_windows_support) ... skipped 'Requires Windows'

----------------------------------------------------------------------
Ran 3 tests in 0.002s

OK (skipped=2)

Creating Subtests

The unittest framework allows you to distinguish between similar tests using the subTest() context manager. For example, say that you have a function that checks whether a given number is even:

def is_even(number):
    return number % 2 == 0

This function uses the modulo operator to check whether the input number is even. Here are some basic tests for the function:

import unittest

from even import is_even

class TestIsEven(unittest.TestCase):
    def test_even_number(self):
        self.assertEqual(is_even(2), True)

    def test_odd_number(self):
        self.assertEqual(is_even(3), False)

    def test_negative_even_number(self):
        self.assertEqual(is_even(-2), True)

    def test_negative_odd_number(self):
        self.assertEqual(is_even(-3), False)

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

These tests take positive and negative numbers and check whether the function’s result is correct by comparing it with the appropriate Boolean value. However, you’ve only tested a couple of input values. If you want to expand the input dataset, then you can use subtests.

To create a subtest, you’ll use the .subTest() method, which returns a context manager that executes the enclosed code block as a subtest. You can use this context manager to provide multiple input values for your tests.

Here’s the above example using subtests to check for multiple numbers:

import unittest

from even import is_even

class TestIsEven(unittest.TestCase):
    def test_even_number(self):
        for number in [2, 4, 6, -8, -10, -12]:
            with self.subTest(number=number):
                self.assertEqual(is_even(number), True)

    def test_odd_number(self):
        for number in [1, 3, 5, -7, -9, -11]:
            with self.subTest(number=number):
                self.assertEqual(is_even(number), False)

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

In this example, you use a for loop that iterates over a list of input values. Then, you use the with statement to manage the context that .subTest() constructs. In that context, you run the assertion with the current number, which works as a subtest. Go ahead and run the test to check the results.

Comparing Values

Comparing the result of a code unit with the expected value is a common way to check whether the unit works okay. The TestCase class defines a rich set of methods that allows you to do this type of check:

You’ve already seen a couple of examples that use .assertEqual(). The .assertNotEqual() method works similarly but with the opposite logic.

To illustrate how the .assertTrue() and .assertFalse() methods work, say that you have the following function:

import math

def is_prime(number):
    if number <= 1:
        return False
    for i in range(2, int(math.sqrt(number)) + 1):
        if number % i == 0:
            return False
    return True

This is a Boolean-valued function, also known as a predicate function, that returns True if the input number is prime and False otherwise. Here’s how you can test this function with unittest:

import unittest

from prime_v1 import is_prime

class TestIsPrime(unittest.TestCase):
    def test_prime_number(self):
        self.assertTrue(is_prime(17))

    def test_non_prime_number(self):
        self.assertFalse(is_prime(10))

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

In this example, you have a TestIsPrime class with two methods. The first method tests is_prime() with a prime number, which must result in True, so you use .assertTrue() for the check. The second method tests is_prime() with a prime number, which must result in False, and you use .assertFalse() in this case.

If you run the script, then you get the following output:

$ python test_prime_v1.py
test_non_prime_number (__main__.TestIsPrime.test_non_prime_number) ... ok
test_prime_number (__main__.TestIsPrime.test_prime_number) ... ok

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

OK

Both tests pass successfully. So, your is_prime() function works okay. In this example, the number of tests is quite small, so you can take advantage of subtests to run is_prime() with several different numbers.

Comparing Objects by Their Identity

TestCase also implements methods that are related to the identity of objects. In Python’s CPython implementation, an object’s identity is the memory address where the object lives. This identity is a unique identifier that distinguishes one object from another.

An object’s identity is a read-only property, which means that you can’t change an object’s identity once you’ve created the object. To check an object’s identity, you’ll use the is and is not operators.

Here are a few assert methods that help you check for an object’s identity:

As you can conclude from this table, these methods are shortcuts for the is and is not operators, as suggested by the method names. With the first two methods, you can compare two objects between them. With the final two methods, you can compare an object against None, which is the Python null value.

Consider the following toy example that checks for the identity of list objects:

import unittest

class TestListIdentity(unittest.TestCase):
    def test_list_aliases(self):
        a = ["Python", "unittest"]
        b = a
        self.assertIs(a, b)

    def test_list_objects(self):
        a = ["Python", "unittest"]
        b = ["Python", "unittest"]
        self.assertIsNot(a, b)

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

In the first test, you create a list containing strings. Then, you derive an alias from the original list. Aliases of an object hold a reference to the same object. So, when you compare them with .assertIs(), your check succeeds because both variables refer to the same object with the same identity.

In the second test, you create two different and independent list objects with the same data. These lists don’t point to the same object in memory, so they have different identities. In this example, you use the .assertIsNot() methods, which should succeed because the compared objects don’t have the same identity.

Go ahead and run the tests:

$ python test_identity.py
test_list_aliases (__main__.TestListIdentity.test_list_aliases) ... ok
test_list_objects (__main__.TestListIdentity.test_list_objects) ... ok

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

OK

When you run the test_identity.py file, you’ll see that both tests pass. Identity tests like this one are useful when you’re testing a function or method that must return cached or singleton objects.

Comparing Collections

Another common need when writing tests is to compare collections, such as lists, tuples, strings, dictionaries, and sets. The TestCase class also has shortcut methods for these types of comparisons. Here’s a summary of those methods:

These methods run equality tests between different collection types. Unlike the methods in the previous section, these compare the values of objects rather than their identities.

Here’s a quick example that showcases several of these methods:

import unittest

class TestCollections(unittest.TestCase):
    def test_sequence_objects(self):
        a = ("H", "e", "l", "l", "o")
        b = "Hello"
        self.assertSequenceEqual(a, b)

    def test_string_objects(self):
        a = "Hello"
        b = "Hello"
        self. assertMultiLineEqual(a, b)

    def test_list_objects(self):
        a = [1, 2, 3, 4, 5]
        b = [1, 2, 3, 4, 5]
        self.assertListEqual(a, b)

    def test_tuple_objects(self):
        a = ("Jane", 25, "New York")
        b = ("Jane", 25, "New York")
        self.assertTupleEqual(a, b)

    def test_dictionary_objects(self):
        a = {"framework": "unittest", "language": "Python"}
        b = {"language": "Python", "framework": "unittest"}
        self.assertDictEqual(a, b)

    def test_set_objects(self):
        a = {1, 2, 4, 3, 5}
        b = {1, 5, 3, 4, 2}
        self.assertSetEqual(a, b)

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

In this example, you have several toy tests. The first test compares a tuple with a string using the .assertSequenceEqual() method because they’re both sequences. The test passes because both sequences contain the same set of characters. Note that you can’t perform this test with the .assertEqual() method, so this is a real specialized use case.

Next, you write similar tests to compare string, list, tuple, dictionary, and set objects. It’s important to highlight that the .assertDictEqual() and .assertSetEqual() methods compare their target objects using the same rules for regular comparisons between dictionaries and sets. They compare the items without considering their order in the collection.

Now, go ahead and run the test_collections.py file from your command line to check the results.

Running Membership Tests

A membership test is a check that allows you to determine whether a given value is or is not in a collection of values. You’ll run these tests with the in and not in operators. Again, the TestCase class has methods for these types of checks:

These two methods provide shortcuts for you to run membership tests on your code. Here’s a quick example of how they work:

import unittest

class TestMembership(unittest.TestCase):
    def test_value_in_collection(self):
        a = 1
        b = [1, 2, 3, 4, 5]
        self.assertIn(a, b)

    def test_value_not_in_collection(self):
        a = 10
        b = [1, 2, 3, 4, 5]
        self.assertNotIn(a, b)

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

In the first test, you check whether a is in the list of values, b. To do this, you use the .assertIn() method. In the second test, you use .assertNotIn() to check whether the a value is not in the target list, b.

Checking for an Object’s Type

Checking the type of the object that a function, method, or callable returns may be another common requirement in testing. For this purpose, the TestClass also has dedicated assert methods:

These two methods are based on the built-in isinstance() function, which you can use to check whether the input object is of a given type.

As an example, say that you have the following class hierarchy:

class Vehicle:
    def __init__(self, make, model):
        self.make = make
        self.model = model

class Car(Vehicle):
    def __init__(self, make, model, max_speed):
        super().__init__(make, model)
        self.max_speed = max_speed

class Truck(Vehicle):
    def __init__(self, make, model, loading_capacity):
        super().__init__(make, model)
        self.loading_capacity = loading_capacity

def vehicle_factory(cls, *args, **kwargs):
    return cls(*args, **kwargs)

In this example, you have a Vehicle class that is at the top of your class hierarchy. Then, you have two concrete classes that inherit from Vehicle and extend it with new attributes. Finally, you have a factory function that you’ll use to create instances of your classes.

To test this code with unittest, you can do something like the following:

import unittest

import vehicles

class TestVehicles(unittest.TestCase):
    def test_vehicle_is_instance_of_car(self):
        car = vehicles.Car("Chevrolet", "Corvette", 194)
        self.assertIsInstance(car, vehicles.Car)

    def test_vehicle_is_instance_of_truck(self):
        truck = vehicles.Truck("Ford", "F-150", 2000)
        self.assertIsInstance(truck, vehicles.Truck)

    def test_vehicles_factory(self):
        car = vehicles.vehicle_factory(
            vehicles.Car,
            make="Toyota",
            model="Corolla",
            max_speed=180,
        )
        self.assertIsInstance(car, vehicles.Vehicle)

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

In the first two tests, you use .assertIsInstance() to check whether the current object is an instance of Car and Truck, respectively. In the final test, you check the vehicle_factory() function. In this case, you use the base class as the comparison type.

Testing for Exceptions

Sometimes, you’ll need to check for exceptions. Yes, sometimes your own code will raise exceptions as part of its behavior. The TestCase class also provides assert methods that allow you to check for exceptions:

The first method allows checking for explicit exceptions without considering the associated error message, and the second method checks for exceptions and considers the associated message using regular expressions.

To illustrate how you can use these methods in your testing code, consider the following reimplementation of your is_prime() function:

from math import sqrt

def is_prime(number):
    if not isinstance(number, int):
        raise TypeError(
            f"integer number expected, got {type(number).__name__}"
        )
    if number < 2:
        raise ValueError(f"integer above 1 expected, got {number}")
    for candidate in range(2, int(sqrt(number)) + 1):
        if number % candidate == 0:
            return False
    return True

In this updated version of is_prime(), you have two conditional statements. The first conditional ensures that the input number is an integer. If that’s not the case, then your code raises a TypeError with an appropriate error message.

The second conditional checks for input numbers greater than 2. If the input number is 1 or lower, then you raise a ValueError. The rest of the code is similar to what you’ve already seen.

How would you write tests to check that the function raises the appropriate exceptions? To do this, you can use the .assertRaises() method:

import unittest

from prime_v2 import is_prime

class TestIsPrime(unittest.TestCase):
    def test_prime_number(self):
        self.assertTrue(is_prime(17))

    def test_non_prime_number(self):
        self.assertFalse(is_prime(10))

    def test_invalid_type_float(self):
        with self.assertRaises(TypeError):
            is_prime(4.5)

    def test_invalid_type_str(self):
        with self.assertRaises(TypeError):
            is_prime("5")

    def test_zero_and_one(self):
        with self.assertRaises(ValueError):
            is_prime(0)
        with self.assertRaises(ValueError):
            is_prime(1)

    def test_negative_number(self):
        with self.assertRaises(ValueError):
            is_prime(-1)

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

The first two tests are familiar. They cover the primary use case of your is_prime() function. The third and fourth tests check for those cases where the function gets an argument of an incorrect type. These tests specifically check for floating-point numbers and strings.

The fifth test checks for those situations where the input is 0 or 1. In those cases, the function must raise a ValueError, so that’s what the assertions catch. Finally, the sixth test checks for negative numbers, which must also raise a ValueError. You can run the test from your command line to check how they work.

The TestCase class also provides some additional assert methods that help you with warnings and logs:

The first two methods allow you to check for warnings, which are a special category of exceptions in Python. A common example of a warning is a DeprecationWarning, which appears when you use deprecated features of the language.

The final two methods allow you to handle logging. These methods return context managers to test whether a message is logged on the logger or one of its children, with at least the given level.

Using Custom Assert Methods

As with many things in Python, you can also create your assert methods to facilitate your test writing process. To do this, you can subclass TestCase and extend the class with new assertion methods.

For example, say that you frequently need to check that all the values in a list are integer numbers. In that case, you can create a test case class like the following:

import unittest

class CustomTestCase(unittest.TestCase):
    def assertAllIntegers(self, values):
        for value in values:
            self.assertIsInstance(
                value,
                int,
            )

This class inherits from TestCase. So, it provides all the assert methods you’ve seen so far. In addition, you extend the functionality of TestCase with a new assert method called .assertAllIntegers(), which takes a list of values and checks whether all the values are integer numbers.

Here’s how you can use this class in practice:

import unittest

class CustomTestCase(unittest.TestCase):
    # ...

class TestIntegerList(CustomTestCase):
    def test_values_are_integers(self):
        integers_list = [1, 2, 3, 4, 5]
        self.assertAllIntegers(integers_list)

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

In this class, you define a test that uses the .assertAllIntegers() method to test a list where all the values are integers. This test must pass successfully. You can experiment with other lists where the values don’t have the correct type to see how the CustomTestCase class works.

Running Tests

With the command-line interface of unittest, you can run tests directly from modules, classes, and from individual test methods. The following sample commands cover those use cases:

$ python -m unittest test_module1 test_module2
$ python -m unittest test_module.TestCase
$ python -m unittest test_module.TestCase.test_method

The first command will run all the tests from test_module1 and test_module2. You can add more modules to the list if you need to. The second command allows you to run all the tests from a given TestCase class. Finally, the last command lets you run a specific test method from a given TestCase class.

As an example, run the following command to execute the tests in your test_age.py file:

$ python -m unittest test_age
.........
----------------------------------------------------------------------
Ran 9 tests in 0.000s

OK

Note that for this command to work, the target module must be available in the import path of your current Python environment. Otherwise, you’ll get an import error. To avoid this issue, you can also run the tests by passing in the module’s file path:

$ python -m unittest test_age.py
.........
----------------------------------------------------------------------
Ran 9 tests in 0.000s

OK

With this command, unittest will find and load the test module from the provided path, which saves you from getting an import error.

Discovering Tests Automatically

The unittest framework supports test discovery. The test loader can inspect each module in a given directory looking for classes derived from TestCase. Then, the loader groups the found classes within a complete test suite.

For example, to discover and run all the tests that you’ve written so far, run the following command in the directory that contains them:

$ python -m unittest discover
...................................
----------------------------------------------------------------------
Ran XX tests in 0.001s

OK

This command locates all tests in the current directory, groups them in a test suite, and finally runs them. You can use the python -m unittest as a shortcut for the above command.

You can use the -s or --start-directory command-line options with the discover subcommand to specify the directory where your tests reside. Other command-line options of discover include:

With these command-line options, you can tweak your test discovery process to fulfill your test automation needs.

Using Command-Line Options

You’ll also find that the unittest CLI supports several command-line options. Here’s a summary of them:

With these command-line options, you can fine-tune how your tests run. In day-to-day test running, the -v option is probably the most commonly used. This option works the same as setting the verbosity argument to 2 when calling the unittest.main() function.

Creating Test Suites With the TestSuite() Constructor

The TestSuite class allows you to create test suites. The class constructor takes the tests argument that must be an iterable of tests or other test suites. So, you can create your test suites like in the following code snippet:

# ...

def make_suite():
    arithmetic_tests = [
        TestArithmeticOperations("test_add"),
        TestArithmeticOperations("test_subtract"),
        TestArithmeticOperations("test_multiply"),
        TestArithmeticOperations("test_divide"),
    ]
    return unittest.TestSuite(tests=arithmetic_tests)

if __name__ == "__main__":
    suite = make_suite()
    runner = unittest.TextTestRunner(verbosity=2)
    runner.run(suite)

In the make_suite() function, you create a list of all tests from the TestArithmeticOperations test case. Then, you create and return a test suite using the TestSuite() constructor with the list of tests as an argument.

To run the suite, you create a TextTestRunner and pass the suite to its .run() method. If you run this file, then you get the following output:

$ python test_calculations.py
test_add (__main__.TestArithmeticOperations.test_add) ... ok
test_subtract (__main__.TestArithmeticOperations.test_subtract) ... ok
test_multiply (__main__.TestArithmeticOperations.test_multiply) ... ok
test_divide (__main__.TestArithmeticOperations.test_divide) ... ok

----------------------------------------------------------------------
Ran 4 tests in 0.000s

OK

This command only runs the four tests in the suite, skipping the rest of the tests in the test_calculations.py file.

Adding Tests to a Suite: .addTest() and .addTests()

You can also use the .addTest() method to add individual tests to an existing suite. To do this, you can do something like the following:

# ...

def make_suite():
    arithmetic_suite = unittest.TestSuite()
    arithmetic_suite.addTest(TestArithmeticOperations("test_add"))
    arithmetic_suite.addTest(TestArithmeticOperations("test_subtract"))
    arithmetic_suite.addTest(TestArithmeticOperations("test_multiply"))
    arithmetic_suite.addTest(TestArithmeticOperations("test_divide"))
    return arithmetic_suite

if __name__ == "__main__":
    suite = make_suite()
    runner = unittest.TextTestRunner(verbosity=2)
    runner.run(suite)

This new version of make_suite() works like the previous section. Instead of using the class constructor to build the test suite, you use the .addTest() method. This approach can be useful when you have an existing test suite, and you need to add more tests to it.

The TestSuite class also has a .addTests() method that you can use to add several tests in one go. This method takes an iterable of test cases, test suites, or a combination of them. Consider the following example that creates a test suite with the statistical tests:

# ...

def make_suite():
    statistical_tests = [
        TestStatisticalOperations("test_mean"),
        TestStatisticalOperations("test_median_odd"),
        TestStatisticalOperations("test_median_even"),
        TestStatisticalOperations("test_median_unsorted"),
        TestStatisticalOperations("test_mode_single"),
        TestStatisticalOperations("test_mode_multiple"),
    ]
    statistical_suite = unittest.TestSuite()
    statistical_suite.addTests(statistical_tests)

    return statistical_suite

if __name__ == "__main__":
    suite = make_suite()
    runner = unittest.TextTestRunner(verbosity=2)
    runner.run(suite)

In this example, you create a list of tests from the TestStatisticalOperations class. Then, you create a TestSuite() instance and add the list of tests using .addTests(). Finally, you return the test suite as usual.

Go ahead and run the file from your command line:

$ python test_calculations.py
test_mean (__main__.TestStatisticalOperations.test_mean) ... ok
test_median_odd (__main__.TestStatisticalOperations.test_median_odd) ... ok
test_median_even (__main__.TestStatisticalOperations.test_median_even) ... ok
test_median_unsorted (__main__.TestStatisticalOperations.test_median_unsorted) ... ok
test_mode_single (__main__.TestStatisticalOperations.test_mode_single) ... ok
test_mode_multiple (__main__.TestStatisticalOperations.test_mode_multiple) ... ok

----------------------------------------------------------------------
Ran 6 tests in 0.000s

OK

This time, the command only runs the statistical tests. As you can see, test suites are a great way to run tests selectively.

Creating Suites With the load_tests() Function

Adding tests to a suite manually can be a tedious task. It can also be error-prone and represent a maintenance burden. Fortunately, unittest has other tools that can help you create test suites quickly.

The load_tests() function is one of these tools. The function is a hook that unittest provides for customizing test loading and suite creation, either for modules or packages of tests.

The function takes three mandatory arguments. Here’s the signature:

def load_tests(loader, standard_tests, pattern):

The loader argument will hold a test loader, which normally is an instance of TestLoader. When you define load_tests() in a module, the standard_tests argument will receive the tests that are loaded from the module by default. When you define the function in a package, the standard_tests will get the test loaded from the package’s __init__.py file. Finally, the pattern argument is a glob pattern that you can use when discovering the tests.

When you call unittest.main() in a test module, the load_tests() function defined in the module gets called automatically, and unittest takes care of passing in the required arguments. This behavior allows you to build a test suite with minimal boilerplate code.

Here’s an example of how to create two test suites. One for the arithmetic tests and another for the statistical tests:

# ...

def load_tests(loader, standard_tests, pattern):
    suite = unittest.TestSuite()
    suite.addTests(loader.loadTestsFromTestCase(TestArithmeticOperations))
    suite.addTests(loader.loadTestsFromTestCase(TestStatisticalOperations))
    return suite

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

After you create the test suite, you use the .addTests() method to add multiple tests in one go. To build the list of tests, you call the .loadTestsFromTestCase() on the loader argument, which is the default test loader.

Note that in this example, you don’t use the standard_tests or the pattern argument. The former is useful in test packages where you’d like to add tests to those found in __init__.py and build a test suite from there. The pattern argument is also for loading tests from packages rather than from modules.

Test Fixtures

As an example of how to use fixtures, say that you have the following implementation of a stack data structure:

class Stack:
    def __init__(self, items=None):
        self.items = list(items) if items is not None else []

    def push(self, item):
        self.items.append(item)

    def pop(self):
        return self.items.pop()

    def __len__(self):
        return len(self.items)

    def __iter__(self):
        return iter(self.items)

    def __reversed__(self):
        return reversed(self.items)

In this module, you define the Stack class that provides the .push() and .pop() methods. The former method appends a new item to the top of the stack, while the latter method removes and returns the item at the top of the stack.

Then, you implement three special methods. The .__len__() method supports the built-in len() function. The .__iter__() method allows you to use Stack instances in for loops. Finally, the .__reversed__() method allows you to support the built-in reversed() function.

Now, say that you need to write a TestCase subclass to test this class. The test class will have to test all the features of your Stack class. In this situation, you have at least two options:

1. You can create a new an independent instance of Stack for every test.
2. You can create a single instance of the class for all the tests.

The first approach sounds tedious and may require a lot of repetitive code. The second approach sounds better, so you decide to go with it. In this situation, you can use a class-level setup fixture to create the instance of Stack and reuse it in every test:

import unittest

from stack import Stack

class TestStack(unittest.TestCase):
    def setUp(self):
        self.stack = Stack()

    def tearDown(self):
        del self.stack

    def test_push(self):
        self.stack.push(1)
        self.assertEqual(self.stack.items, [1])

    def test_pop(self):
        self.stack.push(2)
        item = self.stack.pop()
        self.assertEqual(item, 2)

    def test_len(self):
        self.stack.push(3)
        self.stack.push(4)
        self.assertEqual(len(self.stack), 2)

    def test_iter(self):
        items = [5, 6, 7]
        for item in items:
            self.stack.push(item)
        for stack_item, test_item in zip(self.stack, items):
            self.assertEqual(stack_item, test_item)

    def test_reversed(self):
        items = [5, 6, 7]
        for item in items:
            self.stack.push(item)
        reversed_stack = reversed(self.stack)
        self.assertEqual(list(reversed_stack), [7, 6, 5])

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

In this test module, you create the TestStack to test the different features of your Stack class. The first two methods define the setup and teardown logic, which consist of creating and removing an instance of Stack, respectively.

When you run the tests, unittest automatically calls .setUp() and .tearDown() before and after running each test method. This way, every test method has a fresh instance to work on. In other words, these two methods save you from creating a fresh instance of Stack in each test method.

Class-Level Fixtures

If you use the .setUpClass() and .tearDownClass() class methods, then you can create class-level fixtures. This type of fixture only runs once per test case class. The .setUpClass() method runs before the test methods, and .tearDownClass() runs after all the test methods have run.

This behavior is known as shared fixtures because all the test methods depend on a single setup and teardown run. Note that shared fixtures break test isolation. In other words, the results of a test will depend on previously run tests. So, they should be used with care.

To illustrate how class-level fixtures work, say you have the following Employee class:

class Employee:
    __slots__ = ["name", "age", "job", "salary"]

    def __init__(self, name, age, job, salary):
        self.name = name
        self.age = age
        self.job = job
        self.salary = salary

    def profile(self):
        for attr in self.__slots__:
            print(f"{attr.capitalize()}: {getattr(self, attr)}")
        print()

The class represents an employee of your company. The class has four attributes to store information about each employee. Because you’re planning to create a large number of instances of this class, you use the .__slots__ attribute to reduce the memory footprint.

The employees’ data lives in a CSV file that looks something like the following:

name,age,job,salary
Alice,25,Engineer,50000
Bob,30,Analyst,60000
Jane,35,Manager,80000
John,40,CEO,100000
...

You need to define a function that reads this file and returns a list of employees. This function can look something like the following:

import csv

class Employee:
    # ...

def from_csv_file(file_path):
    with open(file_path) as file:
        reader = csv.DictReader(file)
        employees = []
        for row in reader:
            employees.append(
                Employee(
                    name=row["name"],
                    age=int(row["age"]),
                    job=row["job"],
                    salary=float(row["salary"]),
                )
            )
        return employees

In the from_csv_file() function, you read a CSV file using the DictReader from the csv module. Then, you run a loop to create a list of Employee instances with the read data. Now, you want to write tests for this function:

import os
import unittest
from tempfile import NamedTemporaryFile

from employee import from_csv_file

SAMPLE_CSV = """name,age,job,salary
Alice,25,Engineer,50000
Bob,30,Analyst,60000
Jane,35,Manager,80000
John,40,CEO,100000
"""

class TestFromCsvFile(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.temp_file = NamedTemporaryFile(
            delete=False,
            mode="w",
            newline="",
            suffix=".csv",
        )
        cls.temp_file_name = cls.temp_file.name
        cls.temp_file.write(SAMPLE_CSV)
        cls.temp_file.close()
        cls.employees = from_csv_file(cls.temp_file_name)

    @classmethod
    def tearDownClass(cls):
        os.remove(cls.temp_file_name)

    def test_from_csv_file_total_employees(self):
        self.assertEqual(len(self.employees), 4)

    def test_from_csv_file_employee_attributes(self):
        self.assertEqual(self.employees[0].name, "Alice")
        self.assertEqual(self.employees[0].age, 25)
        self.assertEqual(self.employees[0].job, "Engineer")
        self.assertEqual(self.employees[0].salary, 50000.0)

    def test_from_csv_file_employee_name(self):
        self.assertEqual(self.employees[0].name, "Alice")
        self.assertEqual(self.employees[1].name, "Bob")
        self.assertEqual(self.employees[2].name, "Jane")
        self.assertEqual(self.employees[3].name, "John")

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

In this TestFromCsvFile, you have the .setUpClass() method. In this method, you create a new temporary file using NamedTemporaryFile from the tempfile module.

To populate this file, you use the sample data stored in the SAMPLE_CSV constant. After wiring the sample data to the temporary file, you close the file and call from_csv_file() to create the list of employees. With these actions, you’ve set up everything to test the function. Note that all these actions run only once before all the test methods run.

In the tearDownClass(), you remove the temporary file to clean up your file system and free the acquired resources. This method will run after the test methods have run. Finally, you have three test methods that check different aspects of the generated list of employees.

Module-Level Fixtures

You can also create module-level fixtures. To make these types of fixtures, you need to define the following functions at the module level:

def setUpModule():
    ...

def tearDownModule():
    ...

These fixtures run once per module. The setup fixture runs before all the test cases in the module, and the teardown fixture runs after all the test cases in the module have run.

If an exception happens in the setUpModule() function, then none of the tests in the module will run, and the tearDownModule() won’t run either. If the raised exception is a SkipTest exception, then the module will be reported as skipped instead of an error.

Module-level fixtures are useful when you have several TestCase subclasses in a module, and some of them will benefit from a common setup and teardown logic.

The classic example is a test module with a few test cases that check for database-related functionalities. These tests may need an active connection to the database, which you can create in the setUpModule() function and close in the tearDownModule() function.

A Quick Example: FizzBuzz

Consider the following function that tries to solve the FizzBuzz challenge, where you return "fizz" for numbers divisible by 3, "buzz" for those divisible by 5, and "fizz buzz" for those divisible by both 3 and 5.

Go ahead and create a fizzbuzz.py file and add the following code to it:

def fizzbuzz(number):
    if number % 3 == 0:
        return "fizz"
    elif number % 5 == 0:
        return "buzz"
    elif number % 15 == 0:
        return "fizz buzz"
    else:
        return number

This function works okay for numbers divisible by 3 or 5. However, there is a slight issue with numbers that are divisible by both.

Here’s a unittest test case for your fizzbuzz() function:

import unittest

from fizzbuzz import fizzbuzz

class TestFizzBuzz(unittest.TestCase):
    def test_fizz(self):
        self.assertEqual(fizzbuzz(3), "fizz")

    def test_buzz(self):
        self.assertEqual(fizzbuzz(5), "buzz")

    def test_fizz_buzz(self):
        self.assertEqual(fizzbuzz(15), "fizz buzz")

    def test_neither(self):
        self.assertEqual(fizzbuzz(1), 1)
        self.assertEqual(fizzbuzz(2), 2)
        self.assertEqual(fizzbuzz(4), 4)

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

The first two tests will pass. However, the third test won’t pass because the function has a bug. Go ahead and run the tests so that you can check the output of a failing test:

$ python test_fizzbuzz.py
test_buzz (__main__.TestFizzBuzz.test_buzz) ... ok
test_fizz (__main__.TestFizzBuzz.test_fizz) ... ok
test_fizz_buzz (__main__.TestFizzBuzz.test_fizz_buzz) ... FAIL
test_neither (__main__.TestFizzBuzz.test_neither) ... ok

======================================================================
FAIL: test_fizz_buzz (__main__.TestFizzBuzz.test_fizz_buzz)
----------------------------------------------------------------------
Traceback (most recent call last):
  File ".../test_fizzbuzz.py", line 13, in test_fizz_buzz
    self.assertEqual(fizzbuzz(15), "fizz buzz")
AssertionError: 'fizz' != 'fizz buzz'
- fizz
+ fizz buzz

----------------------------------------------------------------------
Ran 4 tests in 0.000s

FAILED (failures=1)

In this output, the first highlighted line points you to the failing test. The second highlighted line lets you know the exact line where you can find the failing assertion. Finally, the third highlighted line tells you that the actual result and the expected result are unequal.

The failing test tells you that the fizzbuzz() function has an issue with numbers that are divisible by both 3 and 5. So, you need to modify the code to make the test pass:

def fizzbuzz(number):
    if number % 15 == 0:
        return "fizz buzz"
    elif number % 3 == 0:
        return "fizz"
    elif number % 5 == 0:
        return "buzz"
    else:
        return number

In this update, you moved the condition that checks for numbers divisible by 3 and 5 to the beginning of the function. Now you can run the tests again:

$ python test_fizzbuzz.py
test_buzz (__main__.TestFizzBuzz.test_buzz) ... ok
test_fizz (__main__.TestFizzBuzz.test_fizz) ... ok
test_fizz_buzz (__main__.TestFizzBuzz.test_fizz_buzz) ... ok
test_neither (__main__.TestFizzBuzz.test_neither) ... ok

----------------------------------------------------------------------
Ran 4 tests in 0.000s

OK

Your fizzbuzz() function works correctly after the update. This quick example uncovers the essence of code testing, where you write tests to ensure that your code works as expected so you can fix any issue that the tests reveal.

A Test-Driven Example: Rock, Paper, and Scissors

A well-known philosophy around code testing is the test-driven development (TDD) methodology. In TDD, you convert the code’s requirements into test cases before you write the actual code.

For example, say that you want to code a function that takes a number from 0 to 2 and return the strings "rock", "paper", or "scissors". Instead of writing the function right away, you start writing the following tests:

import unittest

from game import rock_paper_scissors

class TestRockPaperScissors(unittest.TestCase):
    def test_out_of_range(self):
        with self.assertRaises(ValueError):
            rock_paper_scissors(42)

    def test_rock(self):
        self.assertEqual(rock_paper_scissors(0), "rock")

    def test_paper(self):
        self.assertEqual(rock_paper_scissors(1), "paper")

    def test_scissors(self):
        self.assertEqual(rock_paper_scissors(2), "scissors")

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

The first test provides a way to check whether the function validates the input argument, which must be a number between 0 and 2. The other tests check the three possible inputs and the expected output for each of them.

Now you can write your function to gradually pass the tests:

def rock_paper_scissors(choice):
    if choice < 0 or choice > 2:
        raise ValueError("choice must be 0, 1, or 2")

In this snippet, you wrote the code to pass the first test. Go ahead and run the tests from your command line:

$ python test_game.py
.FFF
======================================================================
FAIL: test_paper (__main__.TestRockPaperScissors.test_paper)
----------------------------------------------------------------------
Traceback (most recent call last):
  File ".../test_game.py", line 15, in test_paper
    self.assertEqual(rock_paper_scissors(1), "paper")
AssertionError: None != 'paper'

======================================================================
FAIL: test_rock (__main__.TestRockPaperScissors.test_rock)
----------------------------------------------------------------------
Traceback (most recent call last):
  File ".../test_game.py", line 12, in test_rock
    self.assertEqual(rock_paper_scissors(0), "rock")
AssertionError: None != 'rock'

======================================================================
FAIL: test_scissors (__main__.TestRockPaperScissors.test_scissors)
----------------------------------------------------------------------
Traceback (most recent call last):
  File ".../test_game.py", line 18, in test_scissors
    self.assertEqual(rock_paper_scissors(2), "scissors")
AssertionError: None != 'scissors'

----------------------------------------------------------------------
Ran 4 tests in 0.001s

FAILED (failures=3)

In this output, the first line contains the .FFF text. The dot means that you have a passing test, which is the test_out_of_range() test. The Fs mean that you have three failing tests, which are those that deal with the function’s primary goal.

Next, you can write the required code to pass the rest of the tests:

def rock_paper_scissors(choice):
    if choice < 0 or choice > 2:
        raise ValueError("choice must be 0, 1, or 2")

    if choice == 0:
        return "rock"

In this snippet, you add the required code to pass the test_rock() test. Now you can run the tests again:

$ python test_game.py
.F.F
======================================================================
FAIL: test_paper (__main__.TestRockPaperScissors.test_paper)
----------------------------------------------------------------------
Traceback (most recent call last):
  File ".../test_game.py", line 15, in test_paper
    self.assertEqual(rock_paper_scissors(1), "paper")
AssertionError: None != 'paper'

======================================================================
FAIL: test_scissors (__main__.TestRockPaperScissors.test_scissors)
----------------------------------------------------------------------
Traceback (most recent call last):
  File ".../test_game.py", line 18, in test_scissors
    self.assertEqual(rock_paper_scissors(2), "scissors")
AssertionError: None != 'scissors'

----------------------------------------------------------------------
Ran 4 tests in 0.000s

FAILED (failures=2)

This time you have two failing tests, the test_paper() and test_scissors() tests. During the development, you realize that there’s a better way to code the function:

def rock_paper_scissors(choice):
    if choice < 0 or choice > 2:
        raise ValueError("choice must be 0, 1, or 2")

    choices = ["rock", "paper", "scissors"]
    return choices[choice]

In this final implementation of rock_paper_scissors(), you have a list containing the possible return values. Then, you use the input argument to index the list and return the appropriate string.

Go ahead and run your tests again:

$ python test_game.py
....
----------------------------------------------------------------------
Ran 4 tests in 0.000s

OK

Great! Your function passes all the tests. To code this function, you’ve used an incremental development process driven by test cases. At the end of the process, your code works as expected. So, it’s ready to use!