How to Decorate unittest TestCase Class

user_6946463 0
Viewed 68

I'm trying to create my own decorator for a python unit test My goal is to wrap an entire unittest TestCase class and apply certain modifications to it. For some reason, however, the wrapped class seems to be invisible to the test framework.

For these particular tests, I've been trying to use the python built in unittest framework. I've also tried using pytest and following along with this stackoverflow post, but I was still unable to get my decorator to work.

Here is what I've tried:

# test_decorator.py
import unittest
import functools

def wrapper_for_test(test_case):
    @functools.wraps(test_case)
    def wrapper(*args, **kwargs):
        return test_case(*args, **kwargs)

    return wrapper

@wrapper_for_test
class TestClass(unittest.TestCase):
    def test_example(self):
        self.assertTrue(False)

I then run the following:

$ python -m unittest test_decorator.py

----------------------------------------------------------------------
Ran 0 tests in 0.000s

OK

And the test is not found.

But when I comment out @wrapper_for_test, I get:

$ python -m unittest test_decorator.py

FAIL: test_example (test_decorator.TestClass)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "C:\...\test_decorator.py", line 78, in test_example
    self.assertTrue(False)
AssertionError: False is not true

----------------------------------------------------------------------
Ran 1 test in 0.001s

FAILED (failures=1)

The test fails (as expected).

Why isn't the test found when I use the decorator? How can I wrap a unittest TestCase class and have it be found?

pytest python-unittest python
2 Answers

The decorator should wrap a single function, not the class. So the code should be like this:

class TestClass(unittest.TestCase):
    @wrapper_for_test
    def test_example(self):
        self.assertTrue(False)
user_38005520

I discovered that wrapping a class is much different from wrapping a function. Class wrapping is intended to be used for modifying a method of that class. My original intention was to be able to wrap a TestCase class to execute some code before any tests were ran, then run more code once all tests were finished. This could be achieved by modifying the setUpClass() and tearDownClass() methods of the TestCase class:

import unittest

def wrapper_for_test(test_case):
    orig_setUpClass = test_case.setUpClass
    orig_tearDownClass = test_case.tearDownClass

    @classmethod
    def setUpClass(cls):
        # apply modifications
        print("applying modifications")
        # call original setUpClass()
        return orig_setUpClass()

    @classmethod
    def tearDownClass(cls):
        # revert modifications
        print("Cleaning up now that all tests have executed")
        # call original tearDownClass()
        return orig_tearDownClass()

    test_case.setUpClass = setUpClass
    test_case.tearDownClass = tearDownClass
    return test_case

@wrapper_for_test
class TestClass(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        print("setting up the class")

    def test_example(self):
        self.assertEqual(1, 1)
        self.assertTrue(False)

    def test_two(self):
        self.assertEqual(2, 2)

In my case, I ended up going with a modified version of django's TestContextDecorator:

import unittest
import functools


class modify_initial_setup:

    def __init__(self, attr_name=None, kwarg_name=None):
        self.attr_name = attr_name
        self.kwarg_name = kwarg_name

    def enable(self):
        print("Running code before any tests have run")

    def disable(self):
        print("Cleaning up after all tests have finished")

    def __enter__(self):
        return self.enable()

    def __exit__(self, exc_type, exc_value, traceback):
        self.disable()

    def decorate_class(self, cls):
        if issubclass(cls, unittest.TestCase):
            decorated_setUpClass = cls.setUpClass
            decorated_tearDownClass = cls.tearDownClass

            @classmethod
            def setUpClass(inner_cls: unittest.TestCase):
                context = self.enable()
                if self.attr_name:
                    setattr(inner_cls, self.attr_name, context)
                decorated_setUpClass()

            @classmethod
            def tearDownClass(inner_cls: unittest.TestCase):
                context = self.disable()
                if self.attr_name:
                    setattr(inner_cls, self.attr_name, context)
                decorated_tearDownClass()

            cls.setUpClass = setUpClass
            cls.tearDownClass = tearDownClass
            return cls

        raise TypeError(f"Can only decorate subclasses of unittest.TestCase. Not {cls}")

    def decorate_callable(self, func):
        @functools.wraps(func)
        def inner(*args, **kwargs):
            self.enable()
            try:
                return func(*args, **kwargs)
            except Exception as e:
                raise e
            finally:
                self.disable()

        return inner

    def __call__(self, decorated):
        if isinstance(decorated, type):
            return self.decorate_class(decorated)
        elif callable(decorated):
            return self.decorate_callable(decorated)
        raise TypeError(f"Cannot decorate object of type {type(decorated)}")


@modify_initial_setup()
class TestClass(unittest.TestCase):

    # adding this setUpClass() method isn't necessary. it's just here to show that it still works alongside the wrapper.
    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        print("setting up the class")

    def test_example(self):
        self.assertEqual(1, 1)
        self.assertTrue(False)

    def test_two(self):
        self.assertEqual(2, 2)

And this also allows me to decorate individual test functions in the same way if desired. Note that decorating using this class requires the decorator to be called. So decorate like this @wrapper_for_test() instead of @wrapper_for_test.

user_69464630
Related