How YOU can build your own test framework in Python using decorators

TLDR; this article will show how decorators work, some uses for decorators and a bigger example building the start of test framework. It's ~15 min long BUT, there's a very cool project at the end of it - learning how to build your own test framework :)

So, decorators, what are those, and why would I use them? You might have seen something like this before in your Python code @debug or @log right above your function signature like so:

@log
def method():
  pass

Running your method, you might have gotten out some useful information about the method(). If you haven't seen this, that's fine, let me take you on a journey.

References

• Python build real world apps
• Python space exploration NASA
• Python for beginners

Decorators

The @ character denotes something we call decorators; the idea is to decorate a function without the function being aware. It can add capabilities to the function like:

• execution time, how long it takes the function to run from start to finish.
• debugging, information on function name, parameters.

Let's look at a case where we decorate a function:

def decorator(fn):
  print("BEFORE")
  fn()
  print("AFTER")

def log():
  print("[LOG]")

decorator(log)

Running this function results in the following output:

BEFORE
[LOG]
AFTER

Python decorators

Python has a specific way of handling decorators, so much so it lets you use a shorthand @. But first, let's tweak the above code so that it's a good candidate to using the shorthand. Change the code to the following:

def decorator(fn):
  def wrapper()
    print("BEFORE")
    fn()
    print("AFTER")
  return wrapper

def log():
  print("[LOG]")

fn = decorator(log)

What we did was to create an inner function wrapper() in decorator() function and return it. Also, when invoking decorator(), we save the response to fn and invoke that:

fn = decorator(log)

Here's the beautiful part, Python lets us turn the above code:

def log():
  print("[LOG]")

fn = decorator(log)

and turn it into the following:

@decorator
def log():
  print("[LOG]")

log()

Note above how @decorator have replaced the call to decorator() and @decorator is now above the definition of log().

Pause here for a moment and take this in. This is a great start!

Parameters

OK, so great progress so far, we can write decorators, but what happens if we want to use parameters on functions being decorated? Well, let's see what happens:

@decorator
def add(lhs, rhs):
  return lhs + rhs

add(1,1)

You get an error:

TypeError: wrapper() takes 0 positional arguments but 2 were given

So, we're not handling incoming parameters, how to fix?

Well, when you call a function, there are two parameters that Python works with under the hood:

• *args, a list of parameter values.
• **kwargs, a list of tuples where each tuple is the name and value of your parameter.

You can fix the above issue if you send these via your wrapper function into the function being decorated like so:

def wrapper(*args, **kwargs):
    fn(*args, **kwargs)
  return wrapper

Running the code again, it compiles this time, but the output says None, so there seems to be some other problem, where we lose the return value of our decorated function.

How can we solve that?

What we need to do is to ensure the wrapper() function captures the return value from the function being decorated and return that in turn, like so:

def wrapper(*args, **kwargs):
    ret = fn(*args, **kwargs)
    return ret
  return wrapper

Great, so we learned a little more about dealing with parameters and return values with decorators. Let's learn next how we can find out more about the function we decorate.

Instrumentation

We're building up our knowledge steadily so we're able to use it to build a bigger project at the end of the article. Now let's look at how we can find out more about the function we decorate. In Python, you can inspect and find out more about a function by using the help() method. If you only want to find out a functions' name, you can refer to __name__ property like so:

def log():
  print("[LOG]")

print(log.__name__) # prints log

How does this work in a decorator, is there a difference?

Let's try it by adding it to the decorator:

def decorator(fn):
  def wrapper(*args, **kwargs):
    print("calling function: ", fn.__name__)
    ret = fn(*args, **kwargs)
    return ret
  return wrapper

@decorator
def add(lhs, rhs):
  return lhs + rhs

print(add.__name__)
print(add(1,1))

Calling the above code, we get the following output:

wrapper
calling function:  add
2

It seems, add() within the decorator is rightly referred to as add. However, outside of the decorator calling add.__name__, it goes by wrapper, it seems confused, no?

We can fix this using a library called functools that will preserve the function's information like so:

@functools.wraps(fn)

and our full code now looks like:

import functools

def decorator(fn):
  @functools.wraps(fn)
  def wrapper(*args, **kwargs):
    print("calling function: ", fn.__name__)
    ret = fn(*args, **kwargs)
    return ret
  return wrapper

@decorator
def add(lhs, rhs):
  return lhs + rhs

print(add.__name__)
print(add(1,1))

Running this code, we see that it's no longer confused and the add() function is referred to as add in both places.

Ok, we know enough about decorators at this point to be dangerous, let's build something fun, a test framework.

Assignment - build a test framework

DISCLAIMER: I should say there are many test frameworks out there for Python. This is an attempt to put you decoration skills into work, not to add to an existing pile of excellent frameworks, however, you're free to hack away at this project if you find it useful :)

-1- How should it work

Ok, let's start with how it should work, that will drive the architecture and the design.

We want to have the following experience as developers writing tests:

• I just want to add a @test decorator.
• I want to call some type of assertion in my test methods.
• I want to do max one function call to run all tests.
• I want to see tests that succeed and fail, and I want extra info about the failing test so I can correct it
• Nice to have, please add color to the terminal output.

Based on the above, my test file should look something like this:

@test
def testAddShouldSucceed():
  lhs = 1
  rhs = 1
  expected = 2
  expect(add(lhs,rhs), expected, "{lhs} + {rhs} should be equal to {sum}".format(lhs = lhs, rhs = rhs,sum = expected))

@test
def testAddShouldSucceed2():
  lhs = -1
  rhs = 1
  expected = 0
  expect(add(lhs,rhs), expected, "{lhs} + {rhs} should be equal to {sum}".format(lhs = lhs, rhs = rhs,sum = expected))

if __name__ == "__main__":
  run()

-2- Building the parts, decorator(), expect()

Ok, based on the previous part, we need three things:

• @test, a decorator.
• expect() a method that checks if two values are equal.
• run() a method running all functions decorated with @test.

Let's get to work.

1. Create a file decorator.py and give it the following content:

   def test(fn):
      @functools.wraps(fn)
      def wrapper(*args, **kwargs):
        try:
          fn(*args,**kwargs)
          print(". PASS") 
        except TestError as te:
          print(f"{fn.__name__} FAILED with message: '{te.mess}'")
        # fns.append(wrapper)
      return wrapper

The above decorator test() will run test method and print . PASS if the method runs without error. However, if TestError is thrown, you will see FAILED with message . How can a test method crash? Let's implement expect() next.

1. Create a file expect.py and give it the following code:

   from util import TestError

   def expect(lhs, rhs, mess):
     if(lhs == rhs):
       return True
     else:
       raise TestError(mess)

Ok, so return true if all good, and crash if not, got it!. What is TestError?

1. Create a file util.py and give it the following code:

   class TestError(Exception):
     def __init__(self, mess):
       self.mess = mess

Above we're inheriting from Exception and we'e adding a field mess that we store any error message in.

1. Let's revisit decorator.py. Look at the following codeline:

   def test(fn):
      @functools.wraps(fn)
      def wrapper(*args, **kwargs):
        try:
          fn(*args,**kwargs)
          print(". PASS") 
        except TestError as te:
          print(f"{fn.__name__} FAILED with message: '{te.mess}'")
        # fns.append(wrapper)
      return wrapper

Enable the commented out line # fns.append(wrapper). What this line does is to add each test function to a list. We can use that when we implement the run() function. Add run() to the bottom of the file with this code:

   def run():
     for fn in fns:
       fn()

Great, now we just have one thing missing, some colors :)

-3- Give it color

To give it colors let's install a library coloroma.

1. Create a virtual environment

   python3 -m venv test-env

1. Activate the virtual environment

   source test-env/bin/activate

1. Install colorama

   pip install colorama

1. Adjust decorator.py and ensure the code looks like so:

   import functools

   from colorama import Fore
   from colorama import Style
   from util import TestError

   fns = []

   def test(fn):
     @functools.wraps(fn)
       def wrapper(*args, **kwargs):
         try:
           fn(*args,**kwargs)
           print(f"{Fore.GREEN}. PASS{Style.RESET_ALL}")
         except TestError as te:
           print(f"{Fore.RED}{fn.**name**} FAILED with message: '{te.mess}' {Style.RESET_ALL}")
       fns.append(wrapper)
       return wrapper

    def run():
      for fn in fns:
        fn()

1. Create a file demo.py and give it the following content:

   from decorator import test, run
   from expect import expect

   def add(lhs, rhs):
     return lhs + rhs

   @test
   def testAddShouldSucceed():
     lhs = 1
     rhs = 1
     expected = 2
     expect(add(lhs,rhs), expected, "{lhs} + {rhs} should be equal to {sum}".format(lhs = lhs, rhs = rhs,sum = expected))

   @test
   def testAddShouldSucceed2():
     lhs = -1
     rhs = 1
     expected = 0
     expect(add(lhs,rhs), expected, "{lhs} + {rhs} should be equal to {sum}".format(lhs = lhs, rhs = rhs,sum = expected))

   @test
   def testAddShouldFail():
     lhs = 1
     rhs = 2
     expected = 4
     expect(add(lhs,rhs), expected, "{lhs} + {rhs} should be equal to {sum}".format(lhs = lhs, rhs = rhs,sum = expected))

   if **name** == "**main**":
     run()

-4- Demo time

1. Run the program

   python demo.py

Congrats, you've built your first test framework.

Ideas for improvement

So, what's next?

Well, the expect() method is limited in what it can compare, see if you can build other functions capable of comparing lists, objects and so on. Also, how would you test if a method throws a specific exception?

Summary

You've learned about decorators and their use cases. You've also built something useful with them - a test framework. Try it out, change it to your liking. Also, I would recommend reading this excellent post as well, https://realpython.com/primer-on-python-decorators/