Python unittest module

Why to unit test your python source code?

All programmers want their code to be impeccable, but as the saying goes, to err is human, we make mistakes and leave bugs in our source code. Here is where the unit testing comes to our rescue. If you use unit testing from the very beginning of your code development, it will be easier to detect and remove the bugs not only at the beginning stage of development but also post-production.

Unit tests also make us comfortable while refactoring and help to find any bugs left due to an update in the source code. If you are looking for a career in python development, then unit testing your source code is a must for all the big companies. So let us dive into the testing.

Unittest module

Since I have convinced you to use the unit testing with your python source codes, I will illustrate the process in detail. There are various test-runners in python like unittest, nose/nose2, pytest, etc. We will use unittest to test our python source code. The unittest is an inbuilt module and using it is as easy as:-

import unittest

Writing a simple unit test in python

Let us assume, we have the following python code:-

# calculator.py


def add(x, y):
    """Simple addition function"""
    return x + y


def subtract(x, y):
    """Simple subtraction function"""
    return x - y


def multiply(x, y):
    "Simple multiplication function"
    return x * y


def divide(x, y):
    "Simple division function"
    if y == 0:
        raise ValueError("Can not divide a number by 0.")
    return x / y

We will create a new file for the test cases of the code. The general convention is to use either test_filename.py or filename_test.py. We will use test_filename.py. We will keep both the files in the same directory to make the imports and usage relatively easier.

#test_calculator.py

import unittest
import calculator


class TestCalculator(unittest.TestCase):

    def test_add(self):
        self.assertEqual(calculator.add(1, 5), 6)

Code explanation:-

1. We imported the unittest module and calculator.
2. Created a class inheriting from unittest.TestCase().
3. Then we defined our test for the addition function. You must note that the method/function must start with test_. Otherwise, it won’t run. See the example below:-

import unittest
import calculator


class TestCalculator(unittest.TestCase):

    def add_test(self):
        self.assertEqual(calculator.add(1, 5), 6)

# Output


----------------------------------------------------------------------
Ran 0 tests in 0.000s

OK

1. Finally, we have used the assertEqual statement. You can get the list of the assert statements from here and use it as per your case. Few common assert statements are as under:-

Now there are three ways to run the test:-

1. You can run it from the terminal using the code:- python -m unittest test_calculator.py
2. You can also run it from the terminal using the code:- python -m unittest This will automatically detect all the unit tests and run them. The downside of it is that if you have multiple files and tests, it will run all of them.
3. The last and my favorite method is to use the dunders method:

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

Then you can run the test using the code:-python test_calculator.py
The advantage of using this method is that you can run the test from the text-editor as well.
Running the above test will give us the following output:-

.
----------------------------------------------------------------------
Ran 1 test in 0.000s

OK

It has run one test, and the test is OK (passed).

If there is some error in our code, e.g. in multiplication if we have mistyped ‘**‘ instead of ‘*‘. Then running the test will give us the error.

# calculator.py

def multiply(x, y):
    "Simple multiplication function"
    return x ** y

import unittest
import calculator


class TestCalculator(unittest.TestCase):

    def test_mulitply(self):
        self.assertEqual(calculator.multiply(2, 5), 10)


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

The output will be :-

F
======================================================================
FAIL: test_mulitply (__main__.TestCalculator)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/Users/uditvashisht/Desktop/coding/blog/tutorials/indi/test_calculator.py", line 11, in test_mulitply
    self.assertEqual(calculator.multiply(2, 5), 10)
AssertionError: 32 != 10

----------------------------------------------------------------------
Ran 1 test in 0.001s

FAILED (failures=1)

Understanding the test outputs

The test output will have the following parts:-
1. The first line will show the summary result of the execution of all the tests. “.” - test passed and “F” - test failed.
2. If all the tests have passed, the next line will show the number of tests and time taken followed by “OK” in the following line.
3. If any/all the tests fail, the second line will show the name of the test failed, followed by traceback.
4. The error will be raised in the next line.
5. Next line will show the number of tests ran and time taken.
6. The last line will show ‘FAILED’ and the number of failures.

You can also pass your own, AssertionError message as under:-

def test_mulitply(self):
        self.assertEqual(calculator.multiply(2, 5), 10, "Should be 10")


# output 

AssertionError: 32 != 10 : Should be 10

Handling raised error with unittest

In the above example, we have raised a value error in the divide() function. We need to test that divide by zero will raise the error correctly.

def divide(x, y):
    "Simple division function"
    if y == 0:
        raise ValueError("Can not divide a number by 0.")
    return x / y

We will use assertRaises with the context manager and create the following test in our class TestCalculator():-

def test_divide(self):
        with self.assertRaises(ValueError):
            calculator.divide(10, 0)

If we will run the test:-

.
----------------------------------------------------------------------
Ran 1 test in 0.000s

OK

”.” shows that our test has passed. It means that our program will raise a value error when the number is divided by 0.

Using unittest module in a more complex example

We will use the unittest in a more complex example. For that, I downloaded the following code from Corey Schafer’s OOP tutorial.

# employee.py


class Employee:
    """A sample Employee class"""

    def __init__(self, first, last):
        self.first = first
        self.last = last

    @property
    def email(self):
        return f'{self.first}{self.last}@email.com'

    @property
    def fullname(self):
        return f'{self.first.capitalize()} {self.last.capitalize()}'

Now, when we will create an instance of the employee with the first and the last name, it will automatically create the email and full name of the employee. Also, on changing the first or the last name of the employee should change the email and full name. To test the same, we will create the following tests

# test_employee.py

import unittest
from employee import Employee


class TestEmployee(unittest.TestCase):

    def test_email(self):
        emp_1 = Employee('saral', 'gyaan')
        emp_2 = Employee('udit', 'vashisht')
        self.assertEqual(emp_1.email, 'saralgyaan@email.com')
        self.assertEqual(emp_2.email, 'uditvashisht@email.com')

        emp_1.first = "first"
        emp_2.first = "second"

        self.assertEqual(emp_1.email, 'firstgyaan@email.com')
        self.assertEqual(emp_2.email, 'secondvashisht@email.com')

    def test_fullname(self):
        emp_1 = Employee('saral', 'gyaan')
        emp_2 = Employee('udit', 'vashisht')
        self.assertEqual(emp_1.fullname, 'Saral Gyaan')
        self.assertEqual(emp_2.fullname, 'Udit Vashisht')

        emp_1.first = "first"
        emp_2.first = "second"

        self.assertEqual(emp_1.fullname, 'First Gyaan')
        self.assertEqual(emp_2.fullname, 'Second Vashisht')


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

Running this will give the following output:-

..
----------------------------------------------------------------------
Ran 2 tests in 0.000s

OK

Two ‘.’ show that the tests have passed.

Using setUp and tearDown methods in unittest

In the above test, we are creating the instances of the individual for each test and hence violating the ‘DRY’ convention. To overcome this problem, we can use setUp and tearDown methods and change our code as under. For now, we are simply passing the tearDown method, but it is useful when the testing involves creating files, databases etc. and we want to delete them at the end of each test and start with a clean slate. For better a illustration of how it works, we will add print() function in our tests.

# test_employee.py

import unittest
from employee import Employee


class TestEmployee(unittest.TestCase):

    def setUp(self):
        print("Setting up!")
        self.emp_1 = Employee('saral', 'gyaan')
        self.emp_2 = Employee('udit', 'vashisht')

    def tearDown(self):
        print("Tearing down!\n")

    def test_email(self):
        print("Testing email.")
        self.assertEqual(self.emp_1.email, 'saralgyaan@email.com')
        self.assertEqual(self.emp_2.email, 'uditvashisht@email.com')

        self.emp_1.first = "first"
        self.emp_2.first = "second"

        self.assertEqual(self.emp_1.email, 'firstgyaan@email.com')
        self.assertEqual(self.emp_2.email, 'secondvashisht@email.com')

    def test_fullname(self):
        print("Testing Full Name.")
        self.assertEqual(self.emp_1.fullname, 'Saral Gyaan')
        self.assertEqual(self.emp_2.fullname, 'Udit Vashisht')

        self.emp_1.first = "first"
        self.emp_2.first = "second"

        self.assertEqual(self.emp_1.fullname, 'First Gyaan')
        self.assertEqual(self.emp_2.fullname, 'Second Vashisht')


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

Output:-

Setting up!
Testing email.
Tearing down!

Setting up!
Testing Full Name.
Tearing down!

..
----------------------------------------------------------------------
Ran 2 tests in 0.000s

The output depicts that the setUp function ran before every test and tearDown function ran after every test. This could be useful if you are running multiple tests.

In some of the use-cases, it might be useful to have some code run before the whole set of unit tests and something at the end of the unit tests. In such a scenario, you can use two class methods named setUpClass and tearDownClass.

class TestEmployee(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        pass

    @classmethod
    def tearDownClass(cls):
        pass
    ...

Test-driven development

In the above examples, we have developed/written the code and thereafter written the tests for that code. However, many developers will write the tests first and then code. This is called “test-driven development” and is very popular among pro developers.

Assume that you are asked to write a program to find the area of a circle. The easiest function to write is as under:-

# circle.py

from math import pi


def area(radius):
    return pi * r**2

It seems fine, but now try running it as under:-

# circle.py

from math import pi


def area(radius):
    return pi * radius**2


radii = [1, 3, -2, 5 + 2j, True, "radius"]

for radius in radii:
    print(f'Area of the circle is {area(radius)}')

The output of this will be:-

Area of the circle is 3.141592653589793
Area of the circle is 28.274333882308138
Area of the circle is 12.566370614359172
Area of the circle is (65.97344572538566+62.83185307179586j)
Area of the circle is 3.141592653589793
Traceback (most recent call last):
  File "/Users/uditvashisht/Desktop/coding/blog/tutorials/indi/circle.py", line 13, in <module>
    print(f'Area of the circle is {area(radius)}')
  File "/Users/uditvashisht/Desktop/coding/blog/tutorials/indi/circle.py", line 7, in area
    return pi * radius**2
TypeError: unsupported operand type(s) for ** or pow(): 'str' and 'int'

Surprised? So, the above easy looking function has calculated the area for a positive, negative, complex number and boolean radii. Now, let us proceed with test-driven development and start writing the tests:-

# test_circle.py
import unittest
from math import pi
from circle import area


class TestCircle(unittest.TestCase):

    def test_area(self):
        """Test areas when radius >=0"""
        self.assertAlmostEqual(area(2), pi * 2**2)
        self.assertAlmostEqual(area(0), 0)


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

We created a test for radius >=0 and used assertAlmostEqual to assert the value. This test will pass. Now we will integrate the following two cases in our tests:-
1. The function should raise a ValueError for negative radius.
2. The function should raise a TypeError for the radius of a type other than integer and float.

# test_circle.py
import unittest
from math import pi
from circle import area


class TestCircle(unittest.TestCase):

    def test_area(self):
        """Test areas when radius >=0"""
        self.assertAlmostEqual(area(2), pi * 2**2)
        self.assertAlmostEqual(area(0), 0)

    def test_values(self):
        """Raise value error for negative radius"""
        with self.assertRaises(ValueError):
            area(-2)

    def test_types(self):
        """Raise type error for radius other than int or float"""
        with self.assertRaises(ValueError):
            area(True)


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

Running this test will give us the following output:-

.FF
======================================================================
FAIL: test_types (__main__.TestCircle)
Raise type error for radius other than int or float
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/Users/uditvashisht/Desktop/coding/blog/tutorials/indi/test_circle.py", line 22, in test_types
    area(True)
AssertionError: TypeError not raised

======================================================================
FAIL: test_values (__main__.TestCircle)
Raise value error for negative radius
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/Users/uditvashisht/Desktop/coding/blog/tutorials/indi/test_circle.py", line 17, in test_values
    area(-2)
AssertionError: ValueError not raised

----------------------------------------------------------------------
Ran 3 tests in 0.001s

FAILED (failures=2)

So, one of our tests has passed and rest has failed with the following Assertionerrors:-
1. AssertionError: TypeError not raised
2. AssertionError: ValueError not raised

The test output shows to raise specific errors. Now let us change our code as under:-

# circle.py

from math import pi


def area(radius):
    if type(radius) not in [int, float]:
        raise TypeError("Radius must be an integer or float.")

    if radius < 0:
        raise ValueError("Radius can not be negative.")
    return pi * radius**2

Since we have raised the TypeError and ValueError in our code, the tests will pass.

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

OK

Using mocking with unittest for web requests

There are few situations in which we don’t have any control e.g. if we are doing web-scrapping where our function goes to a website and get some information from it. If the website is down, our function will fail but that will also result in our tests getting failed. However, we want our test to fail only when there is some error in our code. We will use mocking to overcome this problem. Let us have a look at the following example:-

# webscrap.py

import requests


def web_scrap():
    response = requests.get('https://www.google.com/')

    if response.ok:
        return response.text
    else:
        return 'Bad Reponse'

The test for checking this code will be as under:-

import unittest
from unittest.mock import patch
import requests
from webscrap import web_scrap


class TestWebScrap(unittest.TestCase):

    def test_webscrap(self):
        with patch('webscrap.requests.get') as m_get:
            m_get.return_value.ok = True
            m_get.return_value.text = 'Success'
            self.assertEqual(web_scrap(), 'Success')

            m_get.return_value.ok = False
            self.assertEqual(web_scrap(), 'Bad Response')


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

1. Here, we have used patch from unittest.mock() and have run it as a context manager.
2. Then, if the response is “Ok”, we have set the text as “Success” and then used assertEqual.
3. If the website is down, then there will get ‘Bad Response’.

The output of the test is:-

.
----------------------------------------------------------------------
Ran 1 test in 0.001s

OK

I will wrap up this tutorial with a hope that you will be comfortable in unit testing your python source code.

In case of any query, you can leave the comment below.

If you liked our tutorial, there are various ways to support us, the easiest is to share this post. You can also follow us on facebook, twitter and youtube.

If you want to support our work. You can do it using Patreon.