Leveraging the power of decorators in Python’s unit-test framework

Vivek Shrivastava
Apr 16 · 7 min read
Image by Peggy und Marco Lachmann-Anke from Pixabay

Unit testing your code is an industry-wide best-practice to ensure that only good quality, properly tested, stable code is shipped to QA, Staging and Production; in that order. Python provides this functionality with its in-built unit-testing framework. One only needs to add the following line to start creating test cases for modules in the code:

import unittest

In this article, I am going to talk about the unit-testing framework, in small part, but the crux of the implementation would be to leverage Decorators while creating the tests to control which aspects to test. I have talked about the Decorators as a concept in one of my previous article, and recommend reading if you wish to get familiar with python concepts like Context managers, Implicit Tuple Unpacking, Magic methods, Generators and Decorators. In this article, however, I will explore how decorators can be used effectively with unit-tests to create flexible, maintainable and scale-able tests.

For the sake of demonstration, I will be creating a unit-test for one of the existing modules in my GitHub repository; the one that shows the use of ‘requests’ module to make API calls. The code for the module as well as the unit-test is available in the same repository on GitHub (don’t worry, there are code samples in the article too).

When we generally create API components within our code, there is a pre-decided schema for the implementation; a common, documented design to aid both back-end as well as front-end development simultaneously. It is prudent to test the API endpoints and schema before its response is tested. Hence, there is a great need to have control over the methods which set the parameters and call the APIs. To aid this, I have added a couple of flags in the original ‘api.py’ file. The first flag is called ‘enableAPICall’ which, as the name suggests, can enable or disable the actual API call. The second flag is called ‘enableStdout’ which controls if print or even logging messages are enabled or disabled. The second flag is not necessary, but I like to have a clean stdout while running tests and avoid messages from the module being tested while running my test cases.

This example can be extended to use-cases where some part of the code need to be tested before proceeding with another set of tests which are time-consuming, resource-intensive or dependent on external factors (like an API response or hardware interfacing). There are testing methodologies to utilize ‘mocked’ objects which are very helpful, say when your code has to interface with a hardware. You wouldn’t want to send a signal to the intended hardware every time the tests are run, but you want to make sure that rest of the program behaves as intended. So, you can use a mock object that behaves like the original one in various scenarios. The concept is good and the python mock library is very powerful. But, every time you write a test, you also configure and control the object’s behavior. For people new to mock, this can be a bit confusing.

Now, sometimes there is no alternative to using mock objects. But, if one is acquainted with Decorators, one can create their own ‘controls’ within their code that serve multiple purposes. For e.g. in the example that I am using, the flag: ‘enableAPICall’ helps me control network calls which can be disabled midway even during execution and also helps me test part of logic during regressions before testing the actual API calls’ responses. So, here is part of my API module with control flags (please find the complete module at my Git repository and also the NOTE regarding the usage of the public APIs from this GitHub repository, powered by Digital Ocean):

import requests
import json

URL_BASE = 'https://api.publicapis.org/'
URL_GET_CATEGORY = 'entries'
URL_GET_RANDOM = 'random'
URL_GET_ALL_CATEGORIES = 'categories'


class API:
    def __init__(self):
        print('Setting up API...')

        # Control parameters
        self.enableAPICall = True
        self.enableStdout = True

        # API parameters
        self.apiURL = None
        self.apiHeader = None
        self.apiPayload = None
        self.response = None

    def getEntry(self, categoryName):
        if self.enableStdout:
            print('\nFetching category:', categoryName)
        self.apiURL = URL_BASE + URL_GET_CATEGORY
        self.apiPayload = {'category': categoryName, 'https': True}

        if self.enableAPICall is True:
            self.response = requests.get(url=self.apiURL, params=self.apiPayload)
            responseJSON = json.loads(self.response.text)
            for key, value in responseJSON.items():
                if self.enableStdout:
                    print(' -', key, '=', value)

In the test module for this particular class, I now, will define a Decorator as shown below:

def disableAPICall(func):
    @wraps(func)
    def wrapper(self, *args, **kwargs):
        flagValue = self.objAPI.enableAPICall
        self.objAPI.enableAPICall = False
        print('Disabling API calls...')
        retVal = func(self, *args, **kwargs)
        print('Re-enabling API calls...')
        self.objAPI.enableAPICall = flagValue
        return retVal
    return wrapper

Now, there are two points of interest in the Decorator defined above. First, it uses a decorator — ‘wraps’. In python, when a function is ‘decorated’, it loses some of its information, namely the function’s name and docstring. In situations where we need to retain the original function’s information, we use ‘wraps’. In short, using ‘wraps’ helps me access the members of the class to which the function being decorated belongs to. Wraps is a part of the ‘functools’ module.

The second point is that once the function is wrapped, I can save the status of its control flag, explicitly disable the network calls and go on towards calling the function. The original status of the control-flag is re-instated after the function has been called. Any arguments and the return values are passed over as in a normal function call.

To apply the above mentioned decorator on a test case method, we just need to add a single line above the function definition as shown below:

@disableAPICall
def testConfig_getRandom(self):
    print('Testing config: getRandom')
    self.objAPI.getRandom()
    self.assertEqual(self.objAPI.apiURL, 'https://api.publicapis.org/random')
    self.assertDictEqual(self.objAPI.apiPayload, {'auth':'null'})

The test method shown above will now call the actual module function sans the actual API call and we can easily assert that the endpoints and schema are set properly.

We achieved the following benefits through this implementation:

  • We created a decorator that can, with a single line addition, can make a test run with the control-flag disabled.
  • The status of the control-flag is retained after the execution of the test. This ensures that decorating one test doesn’t affect other tests in a test suite.
  • There is no effect on the arguments passed or the return value (the exception being if the function returns different values for cases with control-flag disabled).

A more complete picture of the test-cases for the API module is as shown below:

# Python imports
from functools import wraps
import unittest
import json

# Project modules imports
from api import API

# Defines
def disableAPICall(func):
    @wraps(func)
    def wrapper(self, *args, **kwargs):
        flagValue = self.objAPI.enableAPICall
        self.objAPI.enableAPICall = False
        print('Disabling API calls...')
        retVal = func(self, *args, **kwargs)
        print('Re-enabling API calls...')
        self.objAPI.enableAPICall = flagValue
        return retVal
    return wrapper

def disableStdout(func):
    @wraps(func)
    def wrapper(self, *args, **kwargs):
        stdOut = self.objAPI.enableStdout
        self.objAPI.enableStdout = False
        print('Disabling class STDOUTs...')
        retVal = func(self, *args, **kwargs)
        print('Re-enabling class STDOUTs...')
        self.objAPI.enableStdout = stdOut
        return retVal
    return wrapper


class Test_API(unittest.TestCase):
    def setUp(self):
        print('\n--------------------------------')
        self.objAPI = API()

    @disableAPICall
    @disableStdout
    def testConfig_getRandom(self):
        print('Testing config: getRandom')
        self.objAPI.getRandom()
        self.assertEqual(self.objAPI.apiURL, 'https://api.publicapis.org/random')
        self.assertDictEqual(self.objAPI.apiPayload, {'auth':'null'})

    @disableAPICall
    @disableStdout
    def testConfig_getAllCategories(self):
        print('Testing config: getAllCategories')
        self.objAPI.getAllCategories()
        self.assertEqual(self.objAPI.apiURL, 'https://api.publicapis.org/categories')

    @disableAPICall
    @disableStdout
    def testConfig_getEntry(self):
        print('Testing config with category: Business')
        self.objAPI.getEntry('business')
        self.assertEqual(self.objAPI.apiURL, 'https://api.publicapis.org/entries')
        self.assertDictEqual(self.objAPI.apiPayload, {'category':'business', 'https':True})

    def testResponse_getAllCategories(self):
        print('Testing response: getAllCategories')
        categories = self.objAPI.getAllCategories()
        print(' - Validating number of categories returned = 46')
        self.assertEqual(len(categories),46)

    @disableStdout
    def testResponse_getEntry_Business(self):
        print('Testing response for category: Business')
        self.objAPI.getEntry('business')
        responseJSON = json.loads(self.objAPI.response.text)
        print(' - Asserting keys in response JSON...')
        self.assertIn('count', responseJSON.keys())
        self.assertIn('entries', responseJSON.keys())

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

A few pointers about using the unittest class:

  • The class is inherited from unittest.TestCase.
  • The ‘setUp()’ function is called before tests are run. Here, we use it to create an object of the module being tested.
  • There is also a ‘tearDown()’ which takes care of logic to be executed after a test has been run. I have not used it because there is no special logic that needs to run.
  • The ‘tearDown()’ function executes even if the test has failed.

The above code demonstrates tests for checking only the schema for the network call without actually making the call. There are also functions (the last two) which test the actual network call. The only difference between the two is the use of decorators to disable API calls. Intuitively, the tests become focused only on a specific aspect of the module. We test for the parameters in the ‘testConfig_*’ functions while the focus of the tests in the last two are the responses. We will not check the schema in the methods where APIs are called as that has already been tested. This can also be utilized to make hierarchical test-suites which stop further tests if primary test cases fail.

Furthermore, in the example shown above, one can see that I have defined two decorators (for disabling two different control-flags) and have used both in conjunction. This is because decorators can be nested and are called in the order of their usage. For e.g. for the function ‘testConfig_getRandom()’, the output of the test is:

— — — — — — — — — — — — — — — — 
Setting up API…
Disabling API calls…
Disabling class STDOUTs…
Testing config: getAllCategories
Re-enabling class STDOUTs…
Re-enabling API calls…

As the stdout suggests, the decorators are nested in the order that they are specified. The order is:

  1. APIs are disabled
  2. Stdouts are disabled
  3. Function called
  4. Stdouts are enabled
  5. APIs are enabled

I hope this article helps show how to leverage decorators while writing unit-tests in Python. As always, all the code is available on my Git repository: PyProwess; which has a lot more than only this particular example. Do check it out!

In case you have a query, or if you want me to write on a specific feature of the Python programming language, do drop me a message at my e-mail ID given below. Cheers!

Vivek Shrivastava

Written by

Machine Learning & Data Science. E-Mail at platinum012@gmail.com