Writing a Python Unit Test

What it is

A unit test is a small piece of code that calls one function with known input and checks that the result is what you expect.

Python includes the unittest module for this. You write a class that inherits from TestCase, and each method whose name starts with test is one test.

Unit tests run in a CI pipeline on every change, so a broken function is found before it is released.

How it works

You arrange the input, run the code and assert the outcome.

The code under test

# nettools.py
def mask_to_prefix(mask):
    return sum(bin(int(octet)).count("1") for octet in mask.split("."))


def is_valid_vlan(vlan_id):
    return 1 <= vlan_id <= 4094

The test file

# test_nettools.py
import unittest

from nettools import mask_to_prefix, is_valid_vlan


class TestNettools(unittest.TestCase):
    def test_mask_24(self):
        self.assertEqual(mask_to_prefix("255.255.255.0"), 24)

    def test_mask_30(self):
        self.assertEqual(mask_to_prefix("255.255.255.252"), 30)

    def test_vlan_valid(self):
        self.assertTrue(is_valid_vlan(100))

    def test_vlan_invalid(self):
        self.assertFalse(is_valid_vlan(4095))


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

The last two lines let you run the file directly. Each test method is independent, and each should check one idea.

Running the tests

$ python -m unittest test_nettools.py
....
----------------------------------------------------------------------
Ran 4 tests in 0.001s

OK

Each dot is a passing test. A failure shows F, and an unexpected exception shows E. The command python -m unittest discover finds every file named like test_something.py in the folder tree.

Common assert methods

Testing an error looks like this.

def test_bad_mask(self):
    with self.assertRaises(ValueError):
        mask_to_prefix("not.a.mask")

Reading a failure

When an assertion fails, unittest prints the test name, the line and the two values that it compared.

FAIL: test_mask_24 (test_nettools.TestNettools.test_mask_24)
AssertionError: 25 != 24

The first value is what the code returned, and the second is what the test expected. Here the function returned 25 and the test expected 24, so you fix either the code or the test, whichever is wrong. In a pipeline, a failed test stops the build, so the faulty change does not reach production.

Setup and teardown

setUp runs before every test, and tearDown runs after every test. Use them to create and clean up what the tests need, so that they stay independent.

def setUp(self):
    self.devices = ["R1", "R2"]

Mocking an API call

A unit test should not call a real service. The unittest.mock module can replace a function with a stand in that returns a prepared answer.

from unittest.mock import patch, Mock
import device_api


class TestDeviceApi(unittest.TestCase):
    @patch("device_api.requests.get")
    def test_get_names(self, mock_get):
        mock_get.return_value = Mock(
            status_code=200,
            json=lambda: {"items": [{"name": "R1"}, {"name": "R2"}]},
        )
        self.assertEqual(device_api.get_names(), ["R1", "R2"])

The patch decorator swaps requests.get inside the module device_api for the mock while the test runs. The test needs no network and gives the same result every time.

Good tests

• Test one behaviour per method.
• Give each test a name that says what it checks.
• Cover normal input, edge values and bad input.
• Keep tests fast, and keep them independent of each other.

Where pytest fits

pytest is a popular third party runner. It can run unittest tests and also allows plain functions with the assert statement. The exam focuses on the standard unittest module.