Python Unit Testing: Writing Code to Test Your Own Code

Unit testing is a critical practice in software development that involves testing individual units or components of code in isolation to ensure they function correctly.

In this guide, you’ll learn how to write unit tests for your Python code using the Pytest framework.

Steps

Step 1: Install Pytest

First, install the Pytest package:

pip install pytest

Step 2: Organize your Project

Organize your project so that your test files are in a separate directory from your source code. For example:

my_project/
├── src/
│   └── my_module.py
└── tests/
    └── test_my_module.py

Step 3: Write the Module Functions

In your “my_module.py” file, write the actual code you that want to test. For example:

def add_numbers(a, b):
    """
    Add two numbers together.

    :param a: The first number.
    :param b: The second number.
    :return: The sum of the two numbers.
    """
    return a + b


def divide_numbers(a, b):
    """
    Divide one number by another.

    :param a: The numerator.
    :param b: The denominator.
    :return: The result of the division.
    :raises ZeroDivisionError: If the denominator is zero.
    """
    if b == 0:
        raise ZeroDivisionError("Cannot divide by zero")
    return a / b

Step 4: Write your Test Functions

In your “test_my_module.py” file, write functions that start with “test_” to indicate that they are test functions. For example:

import pytest
from src.my_module import add_numbers, divide_numbers


@pytest.mark.parametrize("a, b, expected", [(1, 2, 3), (-1, 1, 0), (-1, -1, -2)])
def test_add_numbers(a, b, expected):
    """
    Test the add_numbers function for different input values.

    :param a: The first number.
    :param b: The second number.
    :param expected: The expected sum of the two numbers.
    """
    assert add_numbers(a, b) == expected


@pytest.mark.parametrize(
    "a, b, expected", [(10, 2, 5), (8, 4, 2), (10, 0, ZeroDivisionError)]
)
def test_divide_numbers(a, b, expected):
    """
    Test the divide_numbers function for different input values.

    :param a: The first number.
    :param b: The second number.
    :param expected: The expected result or exception.
    """
    if expected == ZeroDivisionError:
        with pytest.raises(ZeroDivisionError):
            divide_numbers(a, b)
    else:
        assert divide_numbers(a, b) == expected

Step 5: Run your Test

Finally, run the “test_my_module.py” file to get the test results:

============================= test session starts =============================
collecting ... collected 6 items

test_my_module.py::test_add_numbers[1-2-3] PASSED                        [ '16%']
test_my_module.py::test_add_numbers[-1-1-0] PASSED                       [ '33%']
test_my_module.py::test_add_numbers[-1--1--2] PASSED                     [ '50%']
test_my_module.py::test_divide_numbers[10-2-5] PASSED                    [ '66%']
test_my_module.py::test_divide_numbers[8-4-2] PASSED                     [ '83%']
test_my_module.py::test_divide_numbers[10-0-ZeroDivisionError] PASSED    ['100%']

============================== 6 passed in '0.01s' ==============================