Unit Testing in Python – using unittest module

In order to check the correctness of the written code, we can write verification tests using the unittest module.

Once written tests also help with subsequent changes to the code – if the changes did not cause errors.

We can also use the TDD (Test Driven Development) technique, in which writing tests precedes the creation of the actual program.

I’m going to show a class that tests the correctness of the methods present in the Employee class of the employee.py module. A simple Employee class will contain, in addition to the init () method, an alternative constructor (class method) and methods: set_bonus () – setting a bonus for an employee, get_full_salary () – returns salary (basis + bonus), get_email () – returns employee’s email address.

For this purpose, I’m writing a script containing a class that inherits from unittest.TestCase. Of course, the script has to import the unittest module and the tested module, i.e. the listing of test_employee.py will look like this:

import unittest
from employee import Employee

class EmployeeTestCase(unittest.TestCase):
    def setUp(self):
        self.john_doe = Employee('John', 5000)
        self.jane_doe = Employee.from_full_name('Jane', 'Doe', 6000)

    def test_get_email(self):
        self.assertEqual(self.john_doe.get_email(), 'John@company.com')
        self.assertEqual(self.jane_doe.get_email(), 'Jane.Doe@company.com')

    def test_get_full_salary(self):
        self.assertEqual(self.john_doe.get_full_salary(), 5000)
        self.assertEqual(self.jane_doe.get_full_salary(), 6000)
        self.john_doe.set_bonus(10)
        self.jane_doe.set_bonus(15)
        self.assertEqual(self.john_doe.get_full_salary(), 5500)
        self.assertEqual(self.jane_doe.get_full_salary(), 6900)

To run the above test script correctly, it must be loaded as a module, i.e.

python -m unittest test_employee.py

or attach the call to unittest.main () to the test script, i.e.

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

Based on the written test script, we can write the file employee.py, i.e.

class Employee():
    bonus = 0

    def __init__(self, name, salary):
        self.name = name
        self.salary = salary
    
    def get_email(self):
        pass

    def set_bonus(self, bonus):
        pass

    def get_full_salary(self):
        pass

    @classmethod
    def from_full_name(cls, first_name, last_name, salary):
        name = '{}.{}'.format(first_name, last_name)
        return cls(name, salary)

The stub of the Employee class created in this way, after running the tests, generates 2 failures in the performed tests (despite the presence of several assertions in one test).

Writing the Employee class methods will result in the correct passing of 2 tests, i.e.

class Employee():
    bonus = 0

    def __init__(self, name, salary):
        self.name = name
        self.salary = salary

    def get_email(self):
        return '{}@company.com'.format(self.name)

    def set_bonus(self, bonus):
        self.bonus = bonus

    def get_full_salary(self):
        return self.salary + self.salary * self.bonus / 100

    @classmethod
    def from_full_name(cls, first_name, last_name, salary):
        name = '{}.{}'.format(first_name, last_name)
        return cls(name, salary)

In addition to checking the equality of assertEqual (x, y), we can check various types of assertions, e.g.

  • assertNotEqual(x, y) # checks if x != y
  • assertIsNone(x) # checks if x is None
  • assertTrue(x) # checks if bool (x) is True

The setUp () method from the unittest module is run each time before the next test is performed. Likewise, there is a tearDown () method that executes after each test completes.

There are also setUpClass () and tearDownClass () class methods, which are run before starting and after completing all tests, eg to initiate and then close the database connection.