Practice: Unit Testing with Pytest#
“Unit testing is an important component of the software development process, as it helps ensure your code’s correctness and reliability. In Python, one of the most popular testing frameworks is pytest, due to its simplicity, ease of use, and features.
In this guide, we will go through the basics of unit testing using pytest, with examples to help you get started with writing and running tests for your projects.
What is Unit Testing?#
Unit testing is a software testing technique that focuses on testing individual units of code, usually functions or methods, in isolation.
By testing these individual units of code across your applications, you can ensure that each part of your code works as expected, making it easier to identify and fix bugs, refactor code, and add new features without introducing regressions.
Why use pytest?#
Pytest offers several advantages over other testing frameworks in Python, such as unittest and nose:
Simplified syntax: The syntax for tests is concise and easy-to-understand, making it easier and faster to write them.
Automatic test discovery: Test functions are automatically discovered without the need for explicit registration.
Rich plugin ecosystem: A large number of plugins can extend pytest functionalities and connect it with other tools and services, such as Coverage.py, Django or Elasticsearch.
Getting started with pytest#
Installing pytest#
During this guide we will use Python 3.11, but pytest supports all Python versions above 3.7. You can install pytest using pip:
pip install pytest
Project structure#
We will create a new Python project to go through the different examples. Create a new directory and add the following subdirectories and files:
├── main_project_folder
│ ├── src
│ │ └── __init__.py
│ ├── tests
│ │ └── __init__.py
│ └── main.py
The main.py file will serve as the entry point for your application. In the src directory, we’ll store the source code for your project, and in tests your test code. Also, we include __init__.py files in these directories, which are necessary for pytest to discover the modules, and can remain empty for this tutorial.
Writing your first tests#
Testing a Python function#
Let’s write a simple Python function that we will test. Create a file called math_utils.py in the src folder and add the following code:
def divide(a: int, b: int) -> float:
"""Return the result of a division operation between two numbers.
Args:
a (int): The numerator.
b (int): The denominator.
Returns:
float: The result of the division operation.
Raises:
ZeroDivisionError: If b is zero.
"""
if b == 0:
raise ZeroDivisionError("You can't divide by zero!")
return a / b
Now, we can write some tests for this function. Create a new file called test_math_utils.py in the tests folder and add the following code:
from src.math_utils import divide
def test_divide_positive_numbers() -> None:
"""Test that divide returns the correct result when given two numbers."""
assert divide(1, 2) == 0.5
def test_divide_negative_numbers() -> None:
"""
Test that divide returns the correct result when given a positive and
a negative number.
"""
assert divide(5, -2) == -2.5
assert divide(-2, 5) == -0.4
In this example, we’ve written two test functions called test_divide_positive_numbers and test_divide_negative_numbers that tests the divide function from our math_utils module.
Their results are tested through the assert statement, which checks that the result of the divide function is equal to the expected result. If an assertion fails, pytest will stop and report the failure.
You can add as many assertions as you want in a test function, but it’s generally recommended to keep it small. Also, notice how the test module and function names start with the test_ prefix, this allows pytest to automatically discover them.
Your project structure should now look like this:
├── main_project_folder
│ ├── src
│ │ ├── __init__.py
│ │ └── math_utils.py
│ ├── tests
│ │ ├── __init__.py
│ │ └── test_math_utils.py
│ └── main.py
To run the tests, navigate to the project root directory and execute the following command in your terminal:
pytest
If all tests pass, you can be confident that our divide function works as expected. 🙂
Handling failed tests#
Now, let’s see what happens if a test fails. Change the test_divide_positive_numbers test function to the following:
def test_divide_positive_numbers() -> None:
"""Test that divide returns the correct result when given two integers."""
assert divide(1, 2) == 0.6
Notice that we’ve altered the expected result from 0.5 to 0.6. If we run pytest again, it will report an error. Pytest reports that the test failed, indicating that the expected result was 0.6 and not 0.5. The failed assert statement is also highlighted. By changing the expected result back to 0.5, the test will pass again.
Testing for exceptions#
In our divide function, we raise a ZeroDivisionError if the denominator is zero. Let’s write a test to ensure that this exception is raised correctly. Add the following test function to your test_math_utils.py file:
import pytest
def test_divide_by_zero() -> None:
"""Test that divide raises a ZeroDivisionError when the denominator is zero."""
with pytest.raises(ZeroDivisionError, match="You can't divide by zero!"):
divide(1, 0)
The pytest.raises function checks that the code inside the with block raises the specified exception. You can also provide an optional match argument to verify that the exception message matches a specific pattern.
Writing More Complex Tests#
Use fixtures to reuse test data#
Let’s now look at a more complex example. To avoid the need for instantiating objects in each test you write, you can create fixtures that can be passed as parameters to the test functions. Create a new file called user.py in the src directory and add the following code:
class User:
"""
This class represents a user with attributes ID, name, and age.
Attributes:
id (Integer): The ID of the user, serves as the primary key.
name (String): The name of the user.
age (Integer): The age of the user.
Methods:
greet(): Returns a greeting message that includes the user's name and age.
"""
id: int
name: str
age: int
def __init__(self, id: int, name: str, age: int) -> None:
"""
Initialize a User instance.
Args:
id (int): The ID of the user.
name (str): The name of the user.
age (int): The age of the user.
"""
self.id = id
self.name = name
self.age = age
def greet(self) -> str:
"""
Create a greeting message that includes the user's name and age.
Returns:
str: A greeting message.
"""
return f"Hello, my name is {self.name} and I am {self.age} years old."
To test this User object, create a new file called test_user.py:
import pytest
from src.user import User
@pytest.fixture
def user() -> User:
"""Pytest fixture to create a User instance for testing."""
return User(1, "John Doe", 30)
def test_user_creation(user: User) -> None:
"""Test the creation of a User instance."""
assert user.id == 1
assert user.name == "John Doe"
assert user.age == 30
def test_greet(user: User) -> None:
"""Test the greet method of a User instance."""
greeting: str = user.greet()
assert greeting == "Hello, my name is John Doe and I am 30 years old."
In this example, we’ve created a fixture called user that returns a new User instance. We then passed this fixture as an argument to our test functions, allowing us to reuse the User instance across multiple tests.
To run these tests, simply execute the pytest command again.
Mock to isolate tests#
Your application may need to interact with external services. Testing these interactions can be hard because it involves setting up test databases or APIs. To make testing easier, you can use Python’s unittest.mock library. It allows you to mimic (“mock”) methods, replacing the real connection with a fake one.
Example scenario using sqlalchemy:
# src/user_repository.py (Simplified)
# ... imports ...
# class UserRepository:
# def get_users(self) -> list[User]:
# users = self.session.query(User).all()
# return users
To test this function without connecting to a database, you can utilize unittest.mock.create_autospec:
import pytest
from unittest.mock import create_autospec
from sqlalchemy.orm import Session
from src.user_repository import UserRepository
from src.user import User
@pytest.fixture
def mock_session() -> Session:
"""Pytest fixture to create a mock Session instance."""
session = create_autospec(Session)
return session
def test_get_users(mock_session: Session) -> None:
"""Test the get_users method of the UserRepository class."""
# Create a fake user
fake_user = User(id=1, name="Alice", age=28)
# Mock the Session.query() method to return our fake user
mock_session.query.return_value.all.return_value = [fake_user]
# Create a UserRepository instance with the mocked session
user_repository = UserRepository("postgresql://test:test@test/test")
user_repository.session = mock_session
# Call the get_users method
users = user_repository.get_users()
# Ensure that Session.query() was called with the correct argument
mock_session.query.assert_called_with(User)
# Assert that the method returned our fake user
assert users == [fake_user]
Key Takeaways#
Keep tests simple and focused: Write tests that check one specific aspect of your code.
Use descriptive test function names: E.g.,
test_divide_positive_numbers.Organize your tests: Group related tests in the same module.
Test for edge cases: Cover extreme scenarios and invalid inputs.
Use pytest fixtures: Avoid duplicating setup code.
Mock external dependencies: Isolate your code from external dependencies.
Conclusion#
Unit testing is crucial to check if your code works right and is reliable. Using best practices and pytest’s tools, you can build tests to protect your code from errors and regressions.