Python for test automation: decorators

Python decorators are functions that modify the behavior of other functions or methods. They are applied to functions using the @ symbol followed by the decorator name on the line directly before the function definition. Decorators allow us to add functionality to existing functions without modifying their code, enhancing code reuse, readability, and maintainability. Decorators are commonly used for tasks such as logging, caching, authentication, retrying, or adding additional functionality to functions dynamically. They provide a powerful mechanism for extending and customizing the behavior of functions and methods in Python, facilitating the implementation of cross-cutting concerns and promoting a modular design approach.

Here is a simple

@my_decorator
def my_function():
    ...

Decorators are higher order functions that receive another function as an argument, and return another function. Multiple decorators can be applied on a function by decorating decorators:

@my_first_decorator
@my_second_decorator
def my_function():
    ...

Syntax

The syntax for a basic decorator is as follows:

def my_decorator(func):           # func is the function we decorate
    def wrapper(*args, **kw):     # *args and **kw are they arguments and keyword
        return func(*args, **kw)  # arguments that are passed by the decorator to
                                  # the decorated function
    return wrapper

We can use this decorator on a function:

@my_decorator
def my_function(s: str):
    print(s)


my_function("hello")
>>>hello

Currently our decorator does nothing, but wrap the function in another function and simply runs the function. To see more interesting results, we can add a print call to our decorator as well:

def my_decorator(func):
    def wrapper(*args, **kw):
        print("hello from wrapper")
        return func(*args, **kw)

    return wrapper

Now when we run my_function again, we see:

@my_decorator
def my_function(s: str):
    print(s)


my_function("hello")
>>>hello from wrapper
>>>hello

We can now see that by using a decorator, we can run some code right before a function is executed, and right after a function is executed. For example we could log how long it takes to run a specific function:

import time


def my_logger(func):
    def wrapper(*args, **kw):
        start = time.time()
        result = func(*args, **kw)
        end = time.time()
        print("The function took", end - start, "seconds to complete")
        return result

    return wrapper


@my_logger
def my_function():
    time.sleep(2)
    print("hello")


my_function()
>>>hello
>>>The function took 2.002782106399536 seconds to complete

Timeout decorator

We can construct a decorator that will try to run a function for a certain amount of time, and return the result if it succeeds, and raise an exception if the function call has not succeeded within the given timeout.

def timeout(func):
    def wrapper(*args, *kw):
        end_time = time.time() + 30
        ex = None
        while time.time() < end_time:
            try:
                return func(*args, **kw)
            except Exception as ex:
                time.sleep(2)
        if ex is not None:
            raise ex

    return wrapper

This decorator runs a while loop for 30 seconds, and tries to run the decorated function. In case an exception was raised by the decorated function, we simply sleep for 2 seconds and try again. The while loop exists if the function call does not raise an exception, or when time runs out, at which point the caught exception is raised.

Now the timeout and retry intervals are hard coded into the decorator, but what about if we want to define those values any time we're using the decorator for different functions?

For that we need to pass arguments to the decorator. To do this, we need to utilize two wrapper functions, instead of one. The decorator function timeout receives two optional arguments; timeout and interval corresponding to how long the timeout is and how ofter we run the decorated function.

def timeout(timeout: float = 30.0, interval: float = 2.0):
    def outer_wrapper(func):
        def inner_wrapper(*args, **kw):
            end_time = time.time() + 30
            ex = None
            while time.time() < end_time:
                try:
                    return func(*args, **kw)
                except Exception as ex:
                    time.sleep(2)
            if ex is not None:
                raise ex

        return inner_wrapper

    return outer_wrapper

To test our decorator, we can use it on a function that only raises an exception to see that our decorator actually retries the function for the amount of time we specify, and raises the exception when it does not succeed:

@timeout(timeout=3.0, interval=1.0)
def test():
    raise ValueError("some value error")


test()
>>>... # you will see a traceback of the function call here
>>>ValueError: some value error

As an exercise, you can implement your own retry decorator, which instead of a timeout, has a set number of times the decorator will attempt to run the decorated function.