Writing Unit Tests

Learn to write unit tests for Python functions and classes using the pytest framework.

We'll cover the following

Note: Check the installation guide for pytest before starting this lesson.

The assert statement

Before we write our first unit test, let’s familiarize ourselves with the assert statement in Python.

The assert statement in Python is used for debugging and testing purposes. It allows us to check if a condition is True, and if it isn’t, it raises an exception called AssertionError. This can be helpful in writing unit tests to verify if the expected results match the actual results.

assert condition, message
The assert statement

Here, the condition is the expression that we want to check for being True, and the message is an optional string that we can provide as a custom error message when the assert statement fails.

If the condition in the assert statement is True, the program continues to execute without any issues. However, if the condition is False, the assert statement raises an AssertionError with the optional message as the error message. This allows us to quickly identify and fix issues during testing and debugging.

Python functions

Assuming we already have pytest installed in our environment, the first step is to create a new function to test. Let’s say we have a simple Python function in the addition.py file that adds two numbers/strings:

Press + to interact
def addition(x, y):

    return x + y

Unit test for the function

To write a unit test for the addition function, follow these steps:

  1. Create a new Python file called test.py in the same directory as the original file that contains the addition function.

  2. Import the addition function from the original file into the test.py file.

  3. Write a test function using the pytest framework with a prefix of test in its name. This naming convention helps pytest automatically recognize and execute the test function. Pytest requires the test function names to start with test. Function names that are not of this format aren’t considered test functions by pytest unless we change the default configurations.

  4. Inside the test function, use the assert statement to check if the addition function returns the expected result for different input values. Save the test.py file.

  5. To run the test, navigate to the directory containing the test_file.py file in the terminal and execute the following command:

pytest
The pytest command

Pytest automatically discovers the test files, executes all the test functions inside them, and provides an output indicating if the tests passed or failed. This way, we can quickly verify if our addition function is behaving as expected and catch any issues in our code.

When writing tests with pytest, it’s important to follow a naming convention for our test files and functions to ensure that pytest can automatically discover and execute them. The recommended naming convention is to add test_ as a prefix or suffix to test the file name, and test as a prefix or suffix to the test functions inside the file. We will learn more about test organization and discovery in pytest in upcoming lessons.

If our file name doesn’t follow this convention, we will need to specify the file name in the command:

pytest test.py

Press + to interact
test_file.py
addition.py
from addition import addition
def test_addition_int():
    assert addition(4, 5) == 9
    assert addition(12, 2) == 14
def test_addition_str():
    assert addition('a', 'b') == 'ab'

We can either create two different functions—test_addition_int and test_addition_str—to test the addition of integer and string, or we can test both in one function.

It’s a good practice to make sure our function names are self-explanatory.

Press + to interact
test_file.py
addition.py
from addition import addition
def test_addition():
    assert addition(4, 5) == 9
    assert addition('a', 'b') == 'ab'

The . next to the file name indicates that the test has passed. If the test fails, pytest shows an error message indicating which assertion failed and what the expected and actual results were.

It’s essential to note that unit tests should cover different scenarios and edge cases. In the above example, we test only a few input values. We should also test for invalid input values and ensure that the function raises the appropriate exception.

Below is an example of a test that fails. Run the test and see the output for yourself.

Press + to interact
test_file.py
addition.py
from addition import addition
def test_addition():
    result = addition(4, 5)
    expected = 10
    assert result == expected, f"Expected {expected}, but got {result}"

Unit test for classes

Just like Python functions, we can test methods of Python classes. Let’s dive into the example directly.

Press + to interact
test_file.py
my_class.py
from my_class import MyClass
# method 1
class TestMyClass1:
    def test_addition(self):
        cls_ = MyClass()
        result = cls_.add(3, 4)
        assert result == 7
    def test_subtraction(self):
        cls_ = MyClass()
        result = cls_.subtract(7, 3)
        assert result == 4
# method 2
class TestMyClass2:
    def setup_method(self):
        self.cls_ = MyClass()
    
    def test_addition(self):
        result = self.cls_.add(3, 4)
        assert result == 7
    def test_subtraction(self):
        result = self.cls_.subtract(7, 3)
        assert result == 4

In this test, we define a class called TestMyClass1 that contains two test methods: test_addition and test_subtraction. In each test method, we create an instance of MyClass, call a method on it with some arguments, and then use an assertion to check if the result is what we expect. In TestClass2, instead of creating a new instance of MyClass in each test function, we add setup_method().

In test classes, we cannot have the __init__ constructor because pytest creates a new instance of the test class for each test method. Therefore, any setup or initialization code can be placed in the setup_method() method or in a fixture function. Additionally, defining a constructor can interfere with pytest’s internal mechanisms for test discovery and result reporting.

Note: It is necessary to have Test as a prefix for a test class in Python. If you choose to deviate from this convention, you need to change pytest's configuration. We will learn more about this in the Significance of Test Organization lesson.

Conclusion

In conclusion, writing unit tests is an essential part of software development. It helps ensure that the individual units of a software application function correctly and meet the expected requirements. The pytest framework makes it easy to write and run unit tests in Python. By following the above steps, you can write your first unit test and start testing your software application.