Skip to main content

Be a Problem Solver

Using pytest to Test Your Code

Table of Contents

Have you ever found yourself inheriting legacy code and questioned its functionality after refactoring? Or, have you made changes to your code and wondered if it still works correctly? If you’ve experienced either of these scenarios, it’s time to consider implementing tests for your code. In this article, we will cover the basics of pytest for conducting unit tests.

You can access the code in my repository.

# Preparing Test Cases

Before you start writing your code, it’s a good idea to prepare some test cases to check if your code does what it’s supposed to do. These test cases should be based on what the business or project requires, so you can be sure that your input and output meet those requirements.

For example, if you’re asked to create a function that removes duplicate integers from a list, you should also find out what should happen when the input is an empty list, a list with no duplicates, or a list with only duplicates. Knowing the answers to these questions will help you create the tests you need. Here are some possible answers:

[] -> []
[1, 2, 3] -> [1, 2, 3]
[2, 1, 3, 1] -> [1, 2, 3]
[1, 1, 1, 1] -> [1]

# Conducting Unit Tests for Functions

Now that you have your test cases ready, you can begin testing your code. In the code snippet below, we have a function named test_remove_duplicates(), which takes two arguments: input and expected. The input argument represents the input for the function being tested, while the expected argument represents the anticipated output of the function.

We use the @pytest.mark.parametrize decorator to pass the test cases we’ve prepared earlier as tuples to the test_remove_duplicates() function. It’s important to note that the first parameter of the @pytest.mark.parametrize decorator is a string containing the names of the parameters for the test_remove_duplicates() function, separated by commas.

# test_remove_duplicates.py
import pytest

from remove_duplicates import remove_duplicates

@pytest.mark.parametrize(
    "input, expected",
    [
        ([], []),
        ([1, 2, 3], [1, 2, 3]),
        ([2, 1, 3, 1], [2, 1, 3]),
        ([1, 1, 1, 1], [1]),
    ],
)
def test_remove_duplicates(input, expected):
    assert remove_duplicates(input) == expected

To run the tests, you can use the pytest command in your terminal (make sure you have pytest installed). The output displayed below indicates that none of the four test cases have passed, because we haven’t implemented the remove_duplicates() function yet

$ pytest test_remove_duplicates.py

_______________________ ERROR collecting test_remove_duplicates.py _______________________
ImportError while importing test module 'test_remove_duplicates.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
../.pyenv/versions/3.10.0/lib/python3.10/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
test_remove_duplicates.py:3: in <module>
    from remove_duplicates import remove_duplicates
E   ModuleNotFoundError: No module named 'remove_duplicates'

Assuming we have successfully implemented the remove_duplicates() function in a module called remove_duplicates.py, we can rerun the tests. This time, you should see that all four test cases have passed.

# remove_duplicates.py

def remove_duplicates(input_list: list[int]) -> list[int]:
    return list(dict.fromkeys(input_list))
$ pytest test_remove_duplicates.py

======================== 4 passed in 0.00s =========================

# Conducting Unit Tests for Class Methods

Let’s say we’re presented with a new challenge: removing duplicate integers from a tuple. In this scenario, we decide to create a class to address both tasks. It’s worth noting that creating a class for this purpose may not be strictly necessary; using two functions within a module would suffice. However, for the sake of illustration, we’ll proceed with the class-based approach.

The code snippet below shows the class we created, which has two methods: remove_duplicates_from_list() and remove_duplicates_from_tuple(). Both methods take a list or tuple as input and return a list or tuple with duplicates removed.

# duplicate_remover.py
class DuplicateRemover:
    def remove_duplicates_from_list(self, input_list: list[int]) -> list[int]:
        return list(dict.fromkeys(input_list))

    def remove_duplicates_from_tuple(self, input_tuple: tuple[int]) -> tuple[int]:
        return tuple(dict.fromkeys(input_tuple))

To test the class methods, we need to create an instance of the class and then call the methods on that instance. To avoid creating a new instance for each test case, we can utilize the @pytest.fixture decorator to create a reusable fixture. In the code snippet below, the fixture is named duplicate_remover, and each test function accepts this fixture as an argument, allowing the test function to utilize it

# test_duplicate_remover.py
import pytest

from duplicate_remover import DuplicateRemover

@pytest.fixture(name="duplicate_remover")
def fixture_duplicate_remover():
    return DuplicateRemover()

@pytest.mark.parametrize(
    "input, expected",
    [
        ([], []),
        ([1, 2, 3], [1, 2, 3]),
        ([2, 1, 3, 1], [2, 1, 3]),
        ([1, 1, 1, 1], [1]),
    ],
)
def test_remove_duplicates_from_list(input, expected, duplicate_remover):
    assert duplicate_remover.remove_duplicates_from_list(input) == expected

@pytest.mark.parametrize(
    "input, expected",
    [
        ((), ()),
        ((1, 2, 3), (1, 2, 3)),
        ((2, 1, 3, 1), (2, 1, 3)),
        ((1, 1, 1, 1), tuple([1])),
    ],
)
def test_remove_duplicates_from_tuple(input, expected, duplicate_remover):
    assert duplicate_remover.remove_duplicates_from_tuple(input) == expected

We can run the tests and see that all test cases have passed.

$ pytest test_duplicate_remover.py

======================== 8 passed in 0.01s =========================

# Run Tests for Modules in Directory

In practice, it’s recommended to organize your project with separate directories for tests and application modules to maintain a clean project structure. For instance, the duplicate_remover.py file should reside in a directory named src (or matching your package name). Meanwhile, the test_duplicate_remover.py file should be placed in a directory named tests, and the import statement in test_duplicate_remover.py should be from src.duplicate_remover import DuplicateRemover.

# test_duplicate_remover.py
import test

from src.DuplicateRemover import DuplicateRemover
...

Now, running pytest tests/test_duplicate_remover.py in the current folder won’t work because the module search path (you can confirm this by checking sys.path) begins in the tests folder, and pytest cannot locate the src folder. To address this issue, you can take one of the following approaches:"

  1. python -m pytest tests/test_duplicate_remover.py to run the tests,
  2. add an empty file __init__.py in the tests folder to make it a package.

Both approaches will ensure that the module search path begins in the current folder, allowing pytest to locate both the src and tests folders

# Useful pytest Commandline Flags

pytest offers several command-line flags (used as pytest <FLAG>) to modify the behavior of your tests. Here are some handy ones:

  1. --pdb: This flag allows you to enter the Python debugger (pdb) when a test fails. It’s useful for inspecting the state of your code at the point of failure.
  2. -s: Use this flag to disable the capture of standard output (stdout). It’s helpful when you want to view the output of your code during debugging.
  3. -x or --exitfirst: When this flag is used, pytest exits immediately upon encountering the first error or failed test. It’s particularly useful when you’ve made significant code changes and want to quickly check if your tests still pass.

# Learn More