Testing Decorated Python Functions

Decorators are a great way of adding functionality to a function with minimal impact on the function itself. On top of that, decorator logic can be re-used on other functions that also require the new functionality. For example the following decorator prints a message to standard output every time a function is called.

# plain_main.py
"""Demonstrates a simple decorator."""


def decorator(func):
    """
    A simple decorator that adds printing a message on a function call.

    Args:
        func: The function to decorate.

    Returns:
        The decorated function.
    """
    def inner(*args, **kwargs):
        """Function that is called instead of original function."""
        print('The decorator was called.')
        return func(*args, **kwargs)

    return inner


@decorator
def main():
    print('The main function was called.')


if __name__ == '__main__':
    print('Calling the main function.')
    main()
$ python3 plain_main.py 
Calling the main function.
The decorator was called.
The main function was called.

The drawback of decorators is that the decorator is applied as soon as the interpreter reaches the function definition and it is hard to access the original function without the decorator applied. This might be desirable during testing where testing of the function and the decorator should be separated.

Adding Testing Guard Logic to a Decorator

The solution is to optionally skip the decorator logic if a certain condition is met that is only true during testing. For example, skip the decorator logic if the TESTING environment variable is set.

# check_main.py
"""Demonstrates a decorator with a testing guard."""

import os


def decorator(func):
    """
    A simple decorator that adds printing a message on a function call unless
    the TESTING environment variable is set.

    Args:
        func: The function to decorate.

    Returns:
        The decorated function.
    """
    def inner(*args, **kwargs):
        """Function that is called instead of original function."""
        # Checking for TESTING environment variable
        if os.getenv('TESTING') is not None:
            # Skipping decortor logic
            return func(*args, **kwargs)

        # Running decorator logic
        print('The decorator was called.')
        return func(*args, **kwargs)

    return inner


@decorator
def main():
    print('The main function was called.')


if __name__ == '__main__':
    print('Calling the main function without TESTING set.')
    main()

    print('Calling the main function with TESTING set.')
    os.environ['TESTING'] = ''
    main()
$ python3 check_main.py 
Calling the main function without TESTING set.
The decorator was called.
The main function was called.
Calling the main function with TESTING set.
The main function was called.

As you can see, the decorator logic was executed under normal circumstances (main call on line 39) and was skipped when the TESTING environment variable was set (main call on line 43). The reason was because of the guard statement on line 21 that checks for the TESTING environment variable.

Guard Decorator

You might now say: “great, thank you David. Now I have to rewrite all of my decorator functions!” Ah, but you don’t. If you stay with me through a little more complex decorator code, you won’t have to! The idea is to write a decorator that modifies another decorator’s behaviour.

# guard_main.py
"""Demonstrates a guard decorator."""

import os


def testing_guard(decorator_func):
    """
    Decorator that only applies another decorator if the TESTING environment
    variable is not set.

    Args:
        decorator_func: The decorator function.

    Returns:
        Function that calls a function after applying the decorator if TESTING
        environment variable is not set and calls the plain function if it is set.
    """
    def replacement(original_func):
        """Function that is called instead of original function."""
        def apply_guard(*args, **kwargs):
            """Decides whether to use decorator on function call."""
            if os.getenv('TESTING') is not None:
                return original_func(*args, **kwargs)
            return decorator_func(original_func)(*args, **kwargs)

        return apply_guard
    return replacement


@testing_guard
def decorator(func):
    """
    A simple decorator that adds printing a message on a function call.

    Args:
        func: The function to decorate.

    Returns:
        The decorated function.
    """
    def replacement(*args, **kwargs):
        """Function that is called instead of original function."""
        print('The decorator was called.')
        return func(*args, **kwargs)

    return replacement


@decorator
def main():
    print('The main function was called.')


if __name__ == '__main__':
    print('Calling the main function without TESTING set.')
    main()

    print('Calling the main function with TESTING set.')
    os.environ['TESTING'] = ''
    main()
$ python3 guard_main.py 
Calling the main function without TESTING set.
The decorator was called.
The main function was called.
Calling the main function with TESTING set.
The main function was called.

As you can see, the behaviour of the code is exactly the same but the decorator is in its original form. The reason this works is because the guard decorator gets to intercept each function call and can then decide whether to first apply the decorator or call the plain function on lines 23-25.

On top of not having to re-write decorator functions which you don’t want to execute during testing, you also get to separate the logic that determines whether the decorator is applied from the decorator logic which will reduce the chances of accidentally executing decorator logic as you make changes to the decorator. It also helps you write clear unit tests for both the decorator and the guard decorator.

Working with Pytest

The last consideration is how do you apply this in practice. I would argue that, unless you are testing the decorator or guard decorator, you should always have the TESTING environment variable set. This ensures that you are only testing function logic and not decorator logic. You can achieve this by putting a fixture in your root conftest.py file with autouse set to True.

@pytest.fixture(scope='function', autouse=True)
def set_testing(monkeypatch):
    """Sets the TESTING environment variable."""
    monkeypatch.setenv('TESTING', '')

When you are testing the decorator you would always ensure that the TESTING environment variable is not set. You can achieve that using a fixture that clears the TESTING environment variable that also has autouse set to True as a part of the file that tests the decorator.

@pytest.fixture(scope='function', autouse=True)
def delete_testing(monkeypatch):
    """Deletes the TESTING environment variable."""
    monkeypatch.delenv('TESTING', raising=False)

Finally, for testing the guard decorator, make whether the TESTING environment variable is set part of the tests themselves. Don’t be shy about overriding the functionality of any autouse fixtures as a part of the test function to demonstrate what the intended state of the test is clearly.

def test_guard_decorator_testing_set(monkeypatch):
    """
    GIVEN TESTING environment variable set and ...
    WHEN ...
    THEN ...
    """
    # Setting TESTING environment variable
    monkeypatch.setenv('TESTING', '')

    # Other test code


def test_guard_decorator_testing_not_set(monkeypatch):
    """
    GIVEN TESTING environment variable is not set and ...
    WHEN ...
    THEN ...
    """
    # Setting TESTING environment variable
    monkeypatch.delenv('TESTING', raising=False)

    # Other test code

Thats it! Consider whether environment variables is the best way of indicating that decorator logic should be skipped during testing and also what the best name of the environment variable is for you. I hope this was useful to you!

One thought on “Testing Decorated Python Functions

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Google photo

You are commenting using your Google account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s