如何断言一个方法是用 python unittest 装饰的？

Question

I have a decorator and I want to assert that certain methods in my code are decorated with it.

import functools

def decorator(func):
    def _check_something(*args, **kwargs):
        # some logic in here
        return func(*args, **kwargs)
    return functools.wraps(func)(_check_something)

class MyClass(object):

    @decorator
    def my_method(foo, bar):
        pass

How do I assert with unittest (unitttest2) that my_method has @decorator and no-one removed it, and it was not forgotten?

Answer 1

If for some reason you can't modify the decorator, you could also try checking for some characteristic of a closed variable.

In your example, you know that the original my_method is the only variable closed by the decorator, so you could:

assert (my_method.__closure__ and 
           my_method.__closure__[0].cell_contents.__name__ == my_method.__name__)

Answer 2

You can do that by relying on your decorator to mark the wrapper function with an attribute, that you then assert.

A good practice is to have the decorator set a __wrapped__ attribute pointing to the original function on the returned wrapper.

thus:

def decorator(func):
    @functools.wraps(func)
    def _check_something(*args, **kwargs):
        # some logic in here
        return func(*args, **kwargs)
    _check_something.__wrapped__ = func   # <== add this
    return _check_something

and then, on your test code:

assert getattr(MyClass.my_method, "__wrapped__").__name__ == 'my_method'