How YOU can build your own test framework in Python using decorators
Tags: python, tutorial, programming, testing
Published: 2022-08-31 14:33:36

---
title: How YOU can build your own test framework in Python using decorators
published: true
description: this article will show how decorators work, some uses for decorators and a bigger example building the start of test framework.
tags: python, tutorial, programming, testing
cover_image: https://dev-to-uploads.s3.amazonaws.com/uploads/articles/wh8495qe9asz6osgeubp.png
---

> 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:

```python
@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](https://docs.microsoft.com/en-us/learn/paths/python-language?wt.mc_id=academic-0000-chnoring)
- [Python space exploration NASA](https://docs.microsoft.com/en-us/learn/paths/introduction-python-space-exploration-nasa?wt.mc_id=academic-0000-chnoring)
- [Python for beginners](https://docs.microsoft.com/en-us/learn/paths/beginner-python?wt.mc_id=academic-0000-chnoring)

## 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:

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

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

decorator(log)
```

Running this function results in the following output:

```console
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:

```python
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:

```python
fn = decorator(log)
``` 

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

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

fn = decorator(log)
```

and turn it into the following:

```python
@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:

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

add(1,1)
```

You get an error:

```console
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:

```python
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:

```python
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:

```python
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:

```python
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:

```console
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:

```python
@functools.wraps(fn)
``` 

and our full code now looks like:

```python
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:

```python
@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:

   ```python
   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 **<function> FAILED with message <message>**. How can a test method crash? Let's implement `expect()` next.

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

   ```python
   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:

   ```python
   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:

   ```python
   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:

   ```python
   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

   ```console
   python3 -m venv test-env
   ```

1. Activate the virtual environment

   ```console
   source test-env/bin/activate
   ```

1. Install **colorama**

   ```console
   pip install colorama
   ```

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

   ```python
   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:

   ```python
   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

   ```console
   python demo.py
   ```

   ![colored output](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/643dtissqw5alxbfexgt.png)

> 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/