How To Write a Unit Test in Python: A Simple Guide

Knowing how to write a unit test in Python is critical for developers. Just writing your application code is not enough, writing tests is a must.

Unit tests allow you to test self-contained units of your code independently from each other. Python provides the unittest framework that helps write unit tests following a pre-defined format. To test your code with the unittest framework you create test classes and test methods within each test class.

In this tutorial, you will write unit tests for a simple class that represents a user in a video game.

Let’s learn how to write and run tests in Python!

Why Should You Write Tests in Python?

When you write code you might assume that your code works fine just based on quick manual tests you execute after writing that code.

But, will you run the same tests manually every time you change your code? Unlikely!

To test a piece of code you have to make sure to cover all the different ways that code can be used in your program (e.g. passing different types of values for arguments if you are testing a Python function).

Also, you want to make sure the tests you execute the first time after writing the code are always executed automatically in the future when you change that code. This is why you create a set of tests for your code.

Unit tests are one of the types of tests you can write to test your code. Every unit test covers a specific use case for that code and can be executed every time you modify your code.

Why is it best practice to always implement unit tests and execute unit tests when changing your code?

Executing unit tests helps you make sure not to break the existing logic in your code when you make changes to it (e.g. to provide a new feature). With unit tests, you test code continuously and hence you notice immediately, while writing your code, if you are breaking existing functionality. This allows you to fix errors straight away and keep your Python source code stable.

The Class We Will Write Unit Tests For

The following class represents a user who plays a video game. This class has the following functionalities:

• Activate the user account.
• Check if the user account is active.
• Add points to the user.
• Retrieve points assigned to the user.
• Get the level the user has reached in the game (it depends on the number of points).

The only class attribute is the profile dictionary that stores all the details related to the user.

class User:

    def __init__(self):
        self.profile = {'active': False, 'level': 1, 'points': 0}

    def activate(self):
        self.profile['active'] = True

    def is_active(self):
        return self.profile['active']

    def get_level(self):
        return self.profile['level']

    def get_points(self):
        return self.profile['points']

    def add_points(self, additional_points):
        self.profile['points'] += additional_points

        if self.get_points() > 300:
            self.profile['level'] = 3
        elif self.get_points() > 200:
            self.profile['level'] = 2

Let’s create an instance of this class and run some manual tests to make sure it works as expected.

What is Manual Testing in Python?

Manual testing is the process of testing the functionality of your application by going through use cases one by one.

Think about it as a list of tests you manually run against your application to make sure it behaves as expected. This is also called exploratory testing.

Here is an example…

We will test three different use cases for our class. The first step before doing that is to create an instance of your class:

user1 = User()
print(user1.__dict__)

[output]
{'profile': {'active': False, 'level': 1, 'points': 0}}

As you can see the profile of the user has been initialised correctly.

1st Use Case: User state is active after completing the activation – SUCCESS

user1.activate()
print(user1.is_active())

[output]
True

2nd Use Case: User points are incremented correctly – SUCCESS

user1.add_points(25)
print(user1.get_points())

[output]
25

3rd Use Case: User level changes from 1 to 2 when the number of points is greater than 200 – SUCCESS

print("User total points: {}".format(user1.get_points()))
print("User level: {}".format(user1.get_level()))
user1.add_points(205)
print("User total points: {}".format(user1.get_points()))
print("User level: {}".format(user1.get_level()))

[output]
User total points: 0
User level: 1
User total points: 205
User level: 2

These tests give us some confirmation that the code does what it was designed for.

However, the problem is that you would have to run each single test manually every time the code changes considering that any change could break the existing code.

This is not a great approach, these are just three tests, imagine if you had to run hundreds of tests every time your code changes.

That’s why unit tests are important as a form of automated testing.

How To Write a Unit Test For a Class in Python

Now we will see how to use the Python unittest framework to write the three tests executed in the previous section.

Firstly, let’s assume that the main application code is in the file user.py. You will write your unit tests in a file called test_user.py.

The common naming convention for unit tests: the name of the file used for unit tests simply prepends “test_” to the .py file where the Python code to be tested is.

To use the unittest framework we have to do the following:

• import the unittest module
• create a test class that inherits unittest.TestCase. We will call it TestUser.
• add one method for each test.
• add an entry point to execute the tests from the command line using unittest.main.

Here is a Python unittest example:

import unittest

class TestUser(unittest.TestCase):

    def test_user_activation(self):
        pass

    def test_user_points_update(self):
        pass

    def test_user_level_change(self):
        pass

if __name__ == '__main__':
    unittest.main()

You have created the structure of the test class. Before adding the implementation to each unit test class method let’s try to execute the tests to see what happens.

To run unit tests in Python you can use the following syntax:

$ python test_user.py
...
---------------------------------------------------------------------------
Ran 3 tests in 0.000s

OK

The value of __name__ is checked when you execute the test_user.py file via the command line.

Note: the version of Python we are using in all the examples is Python 3.

How Do You Write a Unit Test in Python?

Now that we have the structure of our test class we can implement each test method.

Unit tests have this name because they test units of your Python code, in this case, the behavior of the methods in the class User.

Each unit test should be designed to verify that the behavior of our class is correct when a specific sequence of events occurs. As part of each unit test, you provide a set of inputs and then verify the output is the same as you expected using the concept of assertions.

In other words, each unit test automates the manual tests we have executed previously.

Technically, you could use the assert statement to verify the value returned by methods of our User class.

In practice, the unittest framework provides its assertion methods. We will use the following in our tests:

• assertEqual
• assertTrue

Let’s start with the first test case…

…actually before doing that we need to be able to see the User class from our test class.

How can we do that?

This is the content of the current directory:

$ ls
test_user.py user.py 

To use the User class in our tests add the following import after the unittest import in test_user.py:

from user import User

And now let’s implement three unit tests.

1st Use Case: User state is active after activation has been completed

def test_user_activation(self):
    user1 = User()
    user1.activate()
    self.assertTrue(user1.is_active())

In this test, we activate the user and then assert that the is_active() method returns True.

2nd Use Case: User points are incremented correctly

def test_user_points_update(self):
    user1 = User()
    user1.add_points(25)
    self.assertEqual(user1.get_points(), 25)

This time instead of using assertTrue we have used assertEqual to verify the number of points assigned to the user.

3rd Use Case: User level changes from 1 to 2 when the number of points is greater than 200

def test_user_level_change(self):
    user1 = User()
    user1.add_points(205)
    self.assertEqual(user1.get_level(), 2)

The implementation of this unit test is similar to the previous one with the only difference being that we are asserting the value of the level for the user.

And now it’s the moment to run our tests…

$ python test_user.py
...
---------------------------------------------------------------------------
Ran 3 tests in 0.000s

OK

All the tests are successful!

An Example of Unit Test Failure

Before completing this tutorial I want to show you what would happen if one of the tests fails.

First of all, let’s assume that there is a typo in the is_active() method:

def is_active(self):
    return self.profile['active_user']

I have replaced the attribute active of the user profile with active_user which doesn’t exist in the profile dictionary.

Now, run the tests again…

$ python test_user.py
E..
===========================================================================
ERROR: test_user_activation (__main__.TestUser)
---------------------------------------------------------------------------
Traceback (most recent call last):
  File "test_user.py", line 9, in test_user_activation
    self.assertTrue(user1.is_active())
  File "/opt/Python/Tutorials/user.py", line 9, in is_active
    return self.profile['active_user']
KeyError: 'active_user'

---------------------------------------------------------------------------
Ran 3 tests in 0.000s

FAILED (errors=1)

In the first line of the test execution you can see:

E..

Each character represents the execution of a test. E indicates an error while a dot indicates a success.

This means that the first test in the test class has failed and the other two are successful.

The output of the test runner also tells us that the error is caused by the assertTrue part of the test_user_activation method.

This helps you identify what’s wrong with the code and fix it.

Conclusion

That was an interesting journey through unit testing in Python.

We have seen how to use unittest, one of the test frameworks provided by Python, to write unit tests for your programs.

Now you have all you need to start writing tests in Python for your application if you haven’t done it before 🙂

Related article: To bring your knowledge of Python testing to the next level have a look at the CodeFatherTech tutorial about creating mocks in Python.