How to Use PyTest for Unit Testing in Python

When you’re writing code in Python, it’s important to make sure that your code works as expected. One of the best ways to do this is by using unit tests, which help you check if small parts (or units) of your code are working correctly.

In this article, we will learn how to write and run effective unit tests in Python using PyTest, one of the most popular testing frameworks for Python.

What are Unit Tests?

Unit tests are small, simple tests that focus on checking a single function or a small piece of code. They help ensure that your code works as expected and can catch bugs early.

Unit tests can be written for different parts of your code, such as functions, methods, and even classes. By writing unit tests, you can test your code without running the entire program.

Why Use PyTest?

PyTest is a popular testing framework for Python that makes it easy to write and run tests.

It’s simple to use and has many useful features like:

• It allows you to write simple and clear test cases.
• It provides advanced features like fixtures, parameterized tests, and plugins.
• It works well with other testing tools and libraries.
• It generates easy-to-read test results and reports.

Setting Up PyTest in Linux

Before we start writing tests, we need to install PyTest. If you don’t have PyTest installed, you can install it using the Python package manager called pip.

pip install pytest

Once PyTest is installed, you’re ready to start writing tests!

Writing Your First Test with PyTest

Let’s start by writing a simple function and then write a test for it.

Step 1: Write a Simple Function

First, let’s create a Python function that we want to test. Let’s say we have a function that adds two numbers:

# add.py
def add(a, b):
    return a + b

This is a simple function that takes two numbers a and b, adds them together, and returns the result.

Step 2: Write a Test for the Function

Now, let’s write a test for the add function. In PyTest, tests are written in separate files, typically named test_*.py to make it easy to identify test files.

Create a new file called test_add.py and write the following test code:

# test_add.py
from add import add

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

Explanation of the above code:

• We import the add function from the add.py file.
• We define a test function called test_add_numbers(). In PyTest, a test function should start with the word test_.
• Inside the test function, we use the assert statement to check if the result of calling the add function matches the expected value. If the condition in the assert statement is True, the test passes; otherwise, it fails.

Step 3: Run the Test

To run the test, open your terminal and navigate to the directory where your test_add.py file is located and then run the following command:

pytest

PyTest will automatically find all the test files (those that start with test_) and run the tests inside them. If everything is working correctly, you should see an output like this:

The dot (.) indicates that the test passed. If there were any issues, PyTest would show an error message.

Writing More Advanced Tests

Now that we know how to write and run a basic test, let’s explore some more advanced features of PyTest.

Testing for Expected Exceptions

Sometimes, you want to test if your code raises the correct exceptions when something goes wrong. You can do this with the pytest.raises() function.

Let’s say we want to test a function that divides two numbers. We want to raise an exception if the second number is zero (to avoid division by zero errors).

Here’s the divide function:

# divide.py
def divide(a, b):
    if b == 0:
        raise ValueError("Cannot divide by zero")
    return a / b

Now, let’s write a test for this function that checks if the ValueError is raised when we try to divide by zero:

# test_divide.py
from divide import divide
import pytest

def test_divide_numbers():
    assert divide(10, 2) == 5
    assert divide(-10, 2) == -5
    assert divide(10, -2) == -5

def test_divide_by_zero():
    with pytest.raises(ValueError):
        divide(10, 0)

Explanation of the code:

• We added a new test function called test_divide_by_zero().
• Inside this function, we use pytest.raises(ValueError) to check if a ValueError is raised when we call the divide function with zero as the second argument.

Run the tests again with the pytest command. If everything is working correctly, you should see this output:

Using Fixtures for Setup and Cleanup

In some cases, you may need to set up certain conditions before running your tests or clean up after the tests are done. PyTest provides fixtures to handle this.

A fixture is a function that you can use to set up or tear down conditions for your tests. Fixtures are often used to create objects or connect to databases that are needed for the tests.

Here’s an example of using a fixture to set up a temporary directory for testing file operations:

# test_file_operations.py
import pytest
import os

@pytest.fixture
def temporary_directory():
    temp_dir = "temp_dir"
    os.mkdir(temp_dir)
    yield temp_dir  # This is where the test will run
    os.rmdir(temp_dir)  # Cleanup after the test

def test_create_file(temporary_directory):
    file_path = os.path.join(temporary_directory, "test_file.txt")
    with open(file_path, "w") as f:
        f.write("Hello, world!")
    
    assert os.path.exists(file_path)

Explanation of the code:

• We define a fixture called temporary_directory that creates a temporary directory before the test and deletes it afterward.
• The test function test_create_file() uses this fixture to create a file in the temporary directory and checks if the file exists.

Run the tests again with the pytest command. PyTest will automatically detect and use the fixture.

Parameterize Your Tests with Pytest

Sometimes, you want to run the same test with different inputs. PyTest allows you to do this easily using parametrize.

Let’s say we want to test our add function with several pairs of numbers. Instead of writing separate test functions for each pair, we can use pytest.mark.parametrize to run the same test with different inputs.

# test_add.py
import pytest
from add import add

@pytest.mark.parametrize("a, b, expected", [
    (2, 3, 5),
    (-1, 1, 0),
    (0, 0, 0),
    (100, 200, 300)
])
def test_add_numbers(a, b, expected):
    assert add(a, b) == expected

Explanation of the code:

• We use the pytest.mark.parametrize decorator to define multiple sets of inputs (a, b, and expected).
• The test function test_add_numbers() will run once for each set of inputs.

Run the tests again with the pytest command, which will run the test four times, once for each set of inputs.

Conclusion

In this article, we’ve learned how to write and run effective unit tests in Python using PyTest to catch bugs early and ensure that your code works as expected.

PyTest makes it easy to write and run these tests, and with its powerful features, you can handle more complex testing needs as you grow in your Python journey.