Unit testing is a type of software testing used to evaluate the smallest unit of code to ensure functionality. These units can be functions, methods, properties or subroutines. In this tutorial, we’re going to teach you how to perform unit testing of your Python code using Python’s built-in unittest module and we’re going to show you several examples.

Testing a Single Function

The unittest module in Python can be used to test Python code on a single function level and with multiple functions. In this section, we’ll focus on testing a single function using the unittest module. Later, we’ll explain how you can test multiple functions.

Writing the Code to be Tested

Let’s first write a simple function take_square() which accepts a number and returns its square. Write the following script in a Python file named arithmetic.py.

def take_square (num1):
    return (num1 * num1)

Running Unit Tests

To run the unit test for testing the functionality of your take_square() function, follow these steps.

Make a new .py file with some name in the same folder as the arithmetic.py you made earlier. At the top of the file, you’re going to first import the unittest module. You don’t need to install the unittest module since it comes pre-built with the default Python installation. You can name this new .py file almost anything you want, but avoid naming it with a name that’s already a standard module name. For example, if you name your file unittest.py, you’ll get the following error:

AttributeError: module 'unittest' has no attribute 'TestCase'

Just change the file name to something like unittesting.py and try again.

Next, you need to import the arithmetic file you created earlier, since the take_square()functions is inside the arithmetic.py file. Again, it’s important you save the following script in a file stored in the same directory as the arithmetic.py file so that the code in the arithmetic.py file can be easily accessed using the simple import command.

import unittest
import arithmetic

class TestArithmetic (unittest.TestCase):

    def test_square(self):
        num = 5
        result = arithmetic.take_square(num)
        self.assertEqual(result, 25)

To actually test the code using unit testing, the script above creates a class which inherits from the TestCase class of the unittest module. In the above script, you create TestArithmetic class which inherits from the unittest.TestCase class. Inside the TestArithmetic class, you define a method test_square(). This is the method which tests the functionality of the take_square() method of the arithmetic.py file.

Inside the test_square() method we declared a variable, num, and assign it a value of 5. The num variable is then passed to the take_square() method of the arithmetic.py file. The value returned by the take_square() function is stored in the result variable.

Next, the self.assertEqual() function from the parent TestCase class tests if the result variable contains 25 (square of 5 is 25). If the self.assertEqual() functions successfully validates the result, the test case is passed.

Finally, to actually run the unit test case, you need to call the main() function from the unittest module. Add the following lines to the end of the script you pasted earlier.

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

If the test case is passed, you’ll see a message like this:

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

OK

From the output, you can see that 1 test ran successfully and the result for that test is OK, which means it’s passed. That means when the unittest module passed 5 to your take_square function, it successfully calculated a result of 25. Your function behaved as expected.

Testing Multiple Functions

You can also unit test multiple Python functions with the unittest module. To do so, you need to update your test class and you’ll have to add a test function for each original function you want to test. Unit tests like this can feel cumbersome but writing good unit tests can save you a tremendous amount of time when debugging since it helps you pinpoint specific units that are failing.

Writing the Code to be Tested

Update your arithmetic.py file by adding the following script.

def take_square (num1):
    return (num1 * num1)


def take_cube (num1):
    return (num1 ** 4)


def take_power(num1, num2):
    return (num1 ** num2)

You should now have three functions in your arithmetic.py file: take_square(), take_cube(), and take_power(). You can see that the take_cube() function is supposed to cube a number, but we intentionally wrote it wrong so it raises the number to the power of 4 instead of 3. We did this deliberately to show you how a failed test case looks when running it through a unit test.

Another thing to notice is our take_power() function accepts two variables. The unittest module isn’t limited to testing basic functions with one input and one output. You can use the module to unit test more complex functions, as well.

Get Our Python Developer Kit for Free

I put together a Python Developer Kit with over 100 pre-built Python scripts covering data structures, Pandas, NumPy, Seaborn, machine learning, file processing, web scraping and a whole lot more - and I want you to have it for free. Enter your email address below and I'll send a copy your way.

Yes, I'll take a free Python Developer Kit

Running Unit Tests for Multiple Functions

Since you want to test three functions, you need to add three corresponding functions to the TestArithmetic class you previously created.

To write unit tests like this, you need to think about the logic of what you intended your function to do and compare the results to what you expected your function to produce. In your mind, you need to ask yourself if I pass this function an input of X, what output should I expect? The output you expect is what you enter in the self.assertEqual line.

Let’s look at the following script. The test_square() function tests the take_square() function, the test_cube() function tests the take_cube() function, and finally the test_power() function tests the take_power() function.

import unittest
import arithmetic

class TestArithmetic (unittest.TestCase):

    def test_square(self):
        num = 5
        result = arithmetic.take_square(num)
        self.assertEqual(result, 25)

    def test_cube(self):
        num = 3
        result = arithmetic.take_cube(num)
        self.assertEqual(result, 27)

    def test_power(self):
        num1 = 2
        num2 = 6

        result = arithmetic.take_power(num1, num2)
        self.assertEqual(result, 64)                

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

The two lines at the end are how we start running our unit tests.

In the output, you can see that one of the test cases, test_cube, has failed. This is because the test case checks if the cube of 3 is 27. However, if you see the original take_cube() function in our arithmetic.py, the functional actually raises the input number to the power of 4 and hence returns 81. Since 81 is not equal to the value specified in the self.assertEqual() function, 27, the test_cube test fails and you’re notified.

F..
======================================================================
FAIL: test_cube (__main__.TestArithmetic)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "<ipython-input-15-eba00b671947>", line 14, in test_cube
    self.assertEqual(result, 27)
AssertionError: 81 != 27

----------------------------------------------------------------------
Ran 3 tests in 0.002s

FAILED (failures=1)

Your unit test identified a deficiency in your function and now you need to fix it. Modify the take_cube() function in the arithmetic.py file to change the 4 to a 3, like this.

def take_cube (num1):
    return (num1 ** 3)

Now the take_cube() function actually returns cube of the integer passed to it. Execute the following script to again run your unit tests.

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

The output below now shows that all the three test cases have been passed.

...
----------------------------------------------------------------------
Ran 3 tests in 0.002s

OK

Get Our Python Developer Kit for Free

I put together a Python Developer Kit with over 100 pre-built Python scripts covering data structures, Pandas, NumPy, Seaborn, machine learning, file processing, web scraping and a whole lot more - and I want you to have it for free. Enter your email address below and I'll send a copy your way.

Yes, I'll take a free Python Developer Kit