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, tests are a must.

Unit tests allow 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 we will write unit tests for a simple class that represents a user in a video game.

Let’s get started!

The Class We Will Write Unit Tests For

The following class represents a user who plays a video game. This class has the follow 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 a 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 our 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.

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

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

[output]
True

Use Case 2: User points are incremented correctly – SUCCESS

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

[output]
25

Use Case 3: User level changes from 1 to 2 when number of points gets higher 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 our code does what it should do.

However, the problem is that we would have to run these tests manually every time the code changes considering that any changes 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 say the main application code is in a file called user.py. We will write our unit tests in a file called test_user.py.

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.

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()

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

$ 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.

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 are called this way because they test units of your code, in this case the behaviour of the methods in the class User.

Each unit test should be designed to verify that the behaviour 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 you expected using the concept of assertions.

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

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

In practice the unittest framework provides its own 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 follow import after the unittest import in test_user.py:

from user import User

And now let’s implement the three unit tests.

Use Case 1: 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.

Use Case 2: 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.

Use Case 3: User level changes from 1 to 2 when number of points gets higher 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 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 that 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 were 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 us identify what’s wrong with our code and fix it.

Conclusion

That was an interesting journey through unit testing in Python.

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