Introduction

Unit testing is a foundational practice in software development involving testing individual components of code to verify that they function correctly. By writing unit tests, developers can identify bugs early, ensure code reliability, and facilitate future changes or enhancements. When speaking about data science — in which algorithms, models, and their code can become complex — unit testing can help maintain the integrity of your code, making it easier to validate results and improve performance.

One of the most popular frameworks for unit testing in Python is pytest, which offers a simple way to write tests, making it easy for developers to ensure their code works as intended. At the heart of pytest unit testing is the assert statement. assert is used to check whether a certain condition is true. If the condition evaluates to false, an AssertionError is raised, indicating that the test has failed.

Let’s look at how to start working with pytest, and how to write and run tests using the assert statement.

Using assert()

Before we start writing tests, we need to set up our environment. This involves installing pytest and creating a directory structure for our project. First, install pytest using pip with the following command:

pip install pytest

Next, create a directory structure for your project. Here’s an example layout:

my_project/
my_module.py
test_my_module.py

In this structure, my_module.py will contain the code we want to test, while test_my_module.py will contain our test cases.

Let’s start by creating a simple function that we will test. Open my_module.py and add the following code:

my_module.py:

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

This function takes two numbers as input and returns their sum. Now, we will create test cases for this function.

In test_my_module.py, we will write our test cases for the add function. First, we need to import pytest and the function we want to test.

test_my_module.py:

from my_module import add

def test_add():
    assert add(2, 3) == 5
    assert add(-1, 1) == 0
    assert add(0, 0) == 0

We have defined a test function called test_add, within which we use the assert statement to check if the output of the add function matches the expected results. If any of these assertions fail, pytest will report a failure.

To run your tests, you can execute the following command from the command line:

pytest

When you run this command, pytest will automatically discover and execute all test functions in files that start with test_. The output will indicate which tests passed and which failed.

=========================================== test session starts ===========================================
platform darwin -- Python 3.12.5, pytest-8.3.3, pluggy-1.5.0
rootdir: /Users/matt/Downloads/pytest-testing
collected 1 item                                                                                                                                  

test_my_module.py .                                                                                 [100%]
============================================ 1 passed in 0.00s ============================================

The dot (.) represents a passed test, while an F would indicate a failure. You will also see detailed information about any failed tests, making it easier to debug your code.

Sometimes, it can be beneficial to provide custom messages for failed assertions, making it easier to understand what went wrong.

assert add(2, 3) == 6, "Expected 5 but got a different result"

In this case, if the assertion fails, pytest will display the custom message alongside the failure.

pytest has built-in assertion rewriting, which makes it easier to write tests without worrying about verbose error messages. For example:

Another useful feature of pytest is the ability to group tests using markers. This can be helpful for organizing tests by categories.

import pytest

@pytest.mark.math
def test_add():
    assert add(2, 3) == 5

Here, the @pytest.mark.math decorator groups the test_add function under the “math” category, which can be useful for running specific groups of tests later on.

Sometimes, you may want to test whether a function raises an exception under certain conditions, which pytest makes straightforward. Here’s an example of how to test for exceptions:

import pytest

def test_add_raises_type_error():
    with pytest.raises(TypeError):
        add("2", "3")

In this test, we expect that calling add with strings will raise a TypeError. If it does, the test will pass; if it doesn’t, the test fails.

When writing tests, you may need to set up some initial state or input data. pytest provides a feature called fixtures to handle this. Here’s a simple example:

import pytest

@pytest.fixture
def input_data():
    return (2, 3)
def test_add_with_fixture(input_data):
    assert add(*input_data) == 5

In this example, we define a fixture called input_data that returns a tuple of numbers. The test function test_add_with_fixture uses this fixture to perform the addition. This approach helps avoid code duplication and keeps tests clean and organized.

Wrapping Up

In this article, we have explored the basics of unit testing using pytest and the assert statement. Practicing writing tests can significantly enhance the reliability of your code and contribute to better software quality, and it is highly recommended that you consider integrating unit tests into your workflow to ensure that your algorithms and models function as expected. For further learning, consider exploring the official pytest documentation and other resources that delve deeper into testing practices in Python.