Python: Unit Testing 101

2 minute read

Set Up

We’ll write unit tests for the following functions

def add(x, y):
   """Add numbers"""
   return x + y

def substract(x, y):
   """Substract numbers"""
   return x - y

def multiply(x, y):
   """Multiply numbers"""

def divide(x, y):
   """Divide numbers"""
   if y == 0:
       raise ValueError('Can not divide by zero!')
   return x / y

Writing the Unit Tests

File name

The file name where our unit tests will reside will be usually called test_<MODULE NAME>.py so in our case the file will be called

Tests class structure

import unittest
import calculator

class TestCalculator(unittest.TestCase):

    def test_<FUNCTION>(self):

    def test_<FUNCTION>(self):

First we import the unittest library which we’ll use for writing and running the tests.
Next, we import the module itself which we aim to test - in our case, the calculator module with the different calculator functions

Then we write the test class. This class inherits from unittest.TestCase which holds multiple useful methods for us to use in our tests classes (more about that later on).
The class itself includes a method for every function we would like to test. In our case, we’ll have a test for every function in the calculator method:

  • test_add
  • test_subtract
  • test_multiply
  • test_divide

‘Add’ Test Function

Let’s implement the ‘add’ test function

import calculator

import unittest

class TestCalculator(unittest.TestCase):

    def test_add(self):
        result = calculator.add(1, 3)
        self.assertEqual(result, 4)

First we use the calculator’s add function to add 1 and 3. We then use one of the methods provided us by the unittest module to assert that both the result of calculator function and the number 4 are equal.
unittest module provide us with multiple assert methods to easily test Python code.

Finally if we’ve used a line like the following self.assertEqual(result, 1) our test method would fail since result is different than 1.

Testing Exceptions

You probably noticed the divide function will raise and exception when you try to divide by 0

if y == 0:
   raise ValueError('Can not divide by zero!')

How do we test such scenario? there are actually at least two ways:

  1. Use assertRaises

self.assertRaises(ValueError, calculator.divide, 5, 0)

  1. Use context manager
with self.assertRaises(ValueError):
    calculator.divide(5, 0)

Running Unit Tests

Run the new test with the following command: python -m unittest
You should see a similar output to the following:

Ran 1 test in 0.000s


You’ll notice that running python will do nothing and the reason is because there is no main entry. So let’s add one

if __name__ == '__main__':

Now running python will run the test similarly to python -m unittest test_calculator

Complete Test Class

import calculator

import unittest

class TestCalculator(unittest.TestCase)

    def test_add(self):
        self.assertEqual(calculator.add(1, 3), 4)
        self.assertEqual(calculator.add(-1, -2), -3)

    def test_substract(self):
        self.assertEqual(calculator.substract(5, 2), 3)
        self.assertEqual(calculator.substract(1, -2), 3)

    def test_multiply(self):
        self.assertEqual(calculator.multiply(2, 3), 6)
        self.assertEqual(calculator.multiply(1, -2), -2)

    def test_divide(self):
        self.assertEqual(calculator.divide(3, 3), 1)
        self.assertEqual(calculator.divide(6, -2), -3)
        with self.assertRaises(ValueError):
            calculator.divide(5, 0)