Assertion Rewriting in Pytest

Pytest plugin for assertion rewriting in Python

Today, we will create a plugin that registers Pytest’s assert rewrite function before any modules are loaded.

Photo by Unseen Studio on Unsplash

Building on yesterday’s article, where we designed our very first Pytest plugin:

Implementing a Pytest Plugin

Pytest Register Assert Rewrite

An important feature of Pytest is its approach to handling the assert keyword. This is a feature we have to understand before we can proceed. As Pytest documentation explains,

One of the main features of pytest is the use of plain assert statements and the detailed introspection of expressions upon assertion failures. This is provided by “assertion rewriting”

Furthermore,

this hook only rewrites test modules themselves (as defined by the python_files configuration option), and any modules which are part of plugins. Any other imported module will not be rewritten and normal assertion behaviour will happen.

For an in-depth understanding, refer to the Pytest Assertion Rewriting section in the Pytest documentation.

Practical Example

To demonstrate how Pytest’s assert rewrite works in practice, we’ll develop a helper module with a comparison function named is_eq. This function will test if two given values are equal:

import logging
from typing import Any


def is_eq(actual: Any, expected: Any, step: str, error_message: str) -> None:
    """Tests if actual is equal to expected.

    :param actual: The actual value.
    :param expected: The expected value.
    :param step: The step.
    :param error_message: The error message if assertion fails.
    """
    logging.info(step)
    assert actual == expected, error_message

Along with this, we’ll design a test case for comparing two dictionaries, utilizing our is_eq function to check their equality:

from tests.assertions import is_eq


class TestDictionary:
    LONG_DICT = {
        "name": "Alice",
        "age": 30,
        "email": "alice@example.com",
        "address": "123, Maple Street, Wonderland",
        "phone": "123-456-7890",
        "occupation": "Engineer",
        "hobbies": ["reading", "cycling", "hiking"],
        "has_pet": True,
        "pet_details": {"pet_name": "Buddy", "pet_type": "Dog", "pet_age": 5},
        "favorite_books": {
            "fiction": "1984",
            "nonfiction": "Sapiens",
            "mystery": "Sherlock Holmes",
            "science fiction": "Dune",
        },
        "languages_spoken": ["English", "Spanish", "French"],
        "education": {
            "undergraduate": "Computer Science",
            "graduate": "Artificial Intelligence",
        },
        "skills": ["Python", "Machine Learning", "Data Analysis"],
        "membership": ["IEEE", "ACM"],
    }

    def test_dictionary(self):
        """Test dictionary."""
        long_dict_copy = self.LONG_DICT.copy()

        # Change name to Bob.
        long_dict_copy["name"] = "Bob"

        is_eq(
            self.LONG_DICT,
            long_dict_copy,
            "The dictionaries shall be equal.",
            "The dictionaries are NOT equal!",
        )

Before and After Pytest Rewrite

When we run these tests without the Pytest assert rewrite, the output is straightforward but not very informative:

$ pytest tests/test_dictionary.py
actual = {'address': '123, Maple Street, Wonderland', 'age': 30, 'education': {'graduate': 'Artificial Intelligence', 'undergraduate': 'Computer Science'}, 'email': 'alice@example.com', ...}
expected = {'address': '123, Maple Street, Wonderland', 'age': 30, 'education': {'graduate': 'Artificial Intelligence', 'undergraduate': 'Computer Science'}, 'email': 'alice@example.com', ...}
step = 'The dictionaries shall be equal.'
error_message = 'The dictionaries are not equal!'

    def is_eq(actual: Any, expected: Any, step: str, error_message: str) -> None:
        """Tests if actual is equal to expected.
    
        :param actual: The actual value.
        :param expected: The expected value.
        :param step: The step.
        :param error_message: The error message if assertion fails.
        """
        logging.info(step)
    
>       assert actual == expected, error_message
E       AssertionError: The dictionaries are not equal!

However, enabling Pytest’s assertion rewrite transforms the output into a more detailed and helpful format. This is clear when comparing the results with and without the rewrite:

$ pytest -p src.py_assert_rewrite tests/test_dictionary.py
actual = {'address': '123, Maple Street, Wonderland', 'age': 30, 'education': {'graduate': 'Artificial Intelligence', 'undergraduate': 'Computer Science'}, 'email': 'alice@example.com', ...}
expected = {'address': '123, Maple Street, Wonderland', 'age': 30, 'education': {'graduate': 'Artificial Intelligence', 'undergraduate': 'Computer Science'}, 'email': 'alice@example.com', ...}
step = 'The dictionaries shall be equal.'
error_message = 'The dictionaries are not equal!'

    def is_eq(actual: Any, expected: Any, step: str, error_message: str) -> None:
        """Tests if actual is equal to expected.
    
        :param actual: The actual value.
        :param expected: The expected value.
        :param step: The step.
        :param error_message: The error message if assertion fails.
        """
        logging.info(step)
    
>       assert actual == expected, error_message
E       AssertionError: The dictionaries are not equal!
E       assert {'address': '...ple.com', ...} == {'address': '...ple.com', ...}
E         Omitting 13 identical items, use -vv to show
E         Differing items:
E         {'name': 'Alice'} != {'name': 'Bob'}
E         Use -v to get more diff

tests/assertions.py:15: AssertionError

The Pytest rewrite significantly improves the readability of the test results, making it easier to spot differences, like {name: 'Alice'} != {name: 'Bob'}.

Pre-Import Registration

In Pytest, a key procedural step is the pre-import registration of assert rewrites. Pytest’s documentation documentation states,

If the helper module also contains assert statements which need to be rewritten it needs to be marked as such, before it gets imported.

So, what happens when a Pytest plugin imports an assertion module itself? How do we ensure the assert keywords within are properly rewritten?

The solution lies in designing a plugin that activates early in the process. By preloading such a plugin, we guarantee that the necessary assert rewrites are in place right from the start, even before the helper module is imported.

Pytest Plugin Setup

We’ll follow the project structure outlined in our previous guide:

py-assert-rewrite/
├─ src/
│  ├─ __init__.py
│  ├─ py_assert_rewrite.py

In this setup, the __init__.py file designates the src folder as a Python package, which can be left blank at this stage.

The critical part of our setup is in the py_assert_rewrite.py file, where we'll include the following code:

import pytest

pytest.register_assert_rewrite("tests.assertions")

This code registers the assert rewrite for our specified module — in this case, tests.assertions. This registration ensures that any assert statements within this module are rewritten.

Running the Plugin

First, ensure that your Python environment recognizes the plugin. This is done by adding the plugin to your PYTHONPATH:

export PYTHONPATH=/path/to/your/py_assert_rewrite:/path/to/your/bin/python

Next, execute your tests using Pytest, ensure to preload your plugin using the -p flag:

$ pytest -p src.py_assert_rewrite tests/test_dictionary.py
============================= test session starts ==============================
...
=================================== FAILURES ===================================
________________________ TestDictionary.test_dictionary ________________________

self = <tests.test_dictionary.TestDictionary object at 0x106ac4e10>

    def test_dictionary(self):
        """Test dictionary."""
        long_dict_copy = self.LONG_DICT.copy()
        long_dict_copy["name"] = "Bob"
    
>       is_eq(
            self.LONG_DICT,
            long_dict_copy,
            "The dictionaries shall be equal.",
            "The dictionaries are not equal!",
        )

tests/test_dictionary.py:38: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

actual = {'address': '123, Maple Street, Wonderland', 'age': 30, 'education': {'graduate': 'Artificial Intelligence', 'undergraduate': 'Computer Science'}, 'email': 'alice@example.com', ...}
expected = {'address': '123, Maple Street, Wonderland', 'age': 30, 'education': {'graduate': 'Artificial Intelligence', 'undergraduate': 'Computer Science'}, 'email': 'alice@example.com', ...}
step = 'The dictionaries shall be equal.'
error_message = 'The dictionaries are not equal!'

    def is_eq(actual: Any, expected: Any, step: str, error_message: str) -> None:
        """Tests if actual is equal to expected.
    
        :param actual: The actual value.
        :param expected: The expected value.
        :param step: The step.
        :param error_message: The error message if assertion fails.
        """
        logging.info(step)
    
>       assert actual == expected, error_message
E       AssertionError: The dictionaries are not equal!
E       assert {'address': '...ple.com', ...} == {'address': '...ple.com', ...}
E         Omitting 13 identical items, use -vv to show
E         Differing items:
E         {'name': 'Alice'} != {'name': 'Bob'}
E         Use -v to get more diff

tests/assertions.py:15: AssertionError
=========================== short test summary info ============================
FAILED tests/test_dictionary.py::TestDictionary::test_dictionary - AssertionE...
============================== 1 failed in 0.02s ===============================

The -p flag here preloads the src.py_assert_rewrite plugin. For further details, refer to the Pytest documentation at Pytest Writing Plugins Guide.

Simplifying Plugin Use with Packaging

Manually setting your Python path each time you open a new terminal window to use the plugin can become tiresome.

A more efficient solution?

Package your plugin and consider publishing it on PyPI. Packaging and publishing your plugin on PyPI enables easy installation with just a simple pip command. For a guide, see:

Packaging and Publishing Your First Python Project

Well Done!

Congratulations on successfully setting up a Pytest plugin that registers assert rewrites before the loading of any helper modules! All the code from our journey is readily available for you to explore and use. Find it on Python Projects on GitHub.

Further Reading

If you want to learn more about programming and, specifically, Python and Java, see the following course:

Introduction to Programming with Python and Java

Note: If you use my links to order, I’ll get a small kickback. So, if you’re inclined to order anything, feel free to click above.