How to mock and test decorator? - python

How to test the following decorator which calls 3rd party library?
import third_party_lib
import functools
class MyDecorator:
def __init__(self, ....):
self.third_party_lib = ThirdPartyLib(....) # will create a 3rd party instance
def __call__(self, ...):
def decorator(f):
#functools.wraps(f)
def wrap(*arg, **kwargs):
result = f(*arg, **kwargs)
# ....
a = self.third_party_lib.send(value=result).get()
# ....
return result
return wrap
return decorator
I need to create an unit test to assert third_party_lib.send() is called if a function is decorated by the decorator. And ideally, also assure the result of the test function is passed to the function.
decorator = MyDecorator(....)
#decorator(....)
def test_func():
ret = ...
return ret # ret should be passed to `third_party_lib.send()`

If you want to verify that the thirdparty function is called correctly, you can mock it and check that the mock is called with the correct parameters. As the ThirdPartyLib intialization shall also be mocked, as mentioned in the comments, you hav to make sure that the docorator is constructed after the mock has been set, for example by constructing it inside the test:
from unittest import mock
#mock.patch('third_party_lib.ThirdPartyLib')
def test_my_decorator(mocked_lib):
decorator = MyDecorator()
#decorator()
def example_func():
return 42
example_func()
mocked_lib.return_value.send.assert_called_once_with(value=42)
If you need the decorated function in more tests, you can wrap it in a function instead:
def decorated_func():
decorator = MyDecorator()
#decorator()
def example_func():
return 42
return example_func
#mock.patch('third_party_lib.ThirdPartyLib')
def test_my_decorator(mocked_lib):
decorated_func()()
mocked_lib.return_value.send.assert_called_once_with(value=42)

Related

Get function or class method name from inside decorator in Python

I have a following decorator.
def allow_disable_in_tests(func):
#functools.wraps(func)
def wrapper(*args, **kwargs):
version = ??? # name of the func of method’s class name
need_to_switch_off_in_tests = cache.get('switch_off_in_tests', version=version)
if settings.IM_IN_TEST_MODE and need_to_switch_off_in_tests:
return None
value = func(*args, **kwargs)
return value
return wrapper
There are 2 types of objects this decorator might accept as a func:
Standalone function.
Method of the class(bound method, static method and class method all possible)
Question is how to get from inside decorator name of the :
function in case of func is standalone function.
method’s class name in case of func is method of the class
This will be used in version .
Decorator should be able to work with both types of objects.
Thanks
You can use __qualname__ and __module__ to derive this information. __qualname__ will describe where the class is defined within a module according to the class or function it was defined in.
However, you are putting test logic in production code, which is a bit of a code smell. You'd be better off using the monkey patching features of your testing framework to patch these functions when running your test suite. For example, with pytest:
import pytest
from functools import wraps
from inspect import signature
class FuncPatch:
def __init__(self, parent, name, retval=None):
self.parent = parent
self.name = name
self.retval = retval
def get_things_to_patch():
import mymodule
return (
FuncPatch(mymodule, 'my_func'),
FuncPatch(mymodule.MyClass, 'method'),
FuncPatch(mymodule.MyClass, 'static'),
FuncPatch(mymodule.MyClass, 'class_', retval='special'),
)
def create_test_function(func, retval, decorator=None):
func = getattr(func, '__func__', func) # unwrap if classmethod or normal method
sig = signature(func)
#wraps(func)
def f(*args, **kwargs):
# check func was called with correct params raises TypeError if wrong
sig.bind(*args, **kwargs)
return retval
if decorator:
f = decorator(f)
return f
#pytest.fixture
def patch_all_the_things(monkeypatch):
for patch in get_things_to_patch():
decorator = None
if (isinstance(patch.parent, type)
and not callable(patch.parent.__dict__[patch.name])
):
# quick hack to detect staticmethod or classmethod
decorator = type(patch.parent.__dict__[patch.name])
to_patch = getattr(patch.parent, patch.name)
func = create_test_function(to_patch, patch.retval, decorator)
monkeypatch.setattr(patch.parent, patch.name, func)
# things to test
def my_func():
return 'my_func'
class MyClass:
#staticmethod
def static():
return 'static'
#classmethod
def class_(cls):
return 'class'
def method(self):
return 'method'
# actual tests
def test_my_func(patch_all_the_things):
assert my_func() is None
def test_my_class(patch_all_the_things):
assert MyClass().method() is None
assert MyClass.method(MyClass()) is None
assert MyClass.static() is None
assert MyClass.class_() == 'special'

python-Decorator function inside Class

I have the below code which works fine if I remove self from methods
class tests:
def __init__(self):
pass
def func(self,a):
def wrapp(x):
y=x+2
return a(y)
return wrapp
#func
def func1(self,b):
return b
print (tests.func1(10))
I believe decorator function are functions that return another function. Will that not work inside class? Ignore the indentation error as I am not achievable when I paste the code here..
Please help me how I can achieve this scenario inside class..
You can just declare your decorator outside of the class. Also, when you are decorating a class method, it seems you need to pass the self variable from the wrapper to the decorated function (changed the names for more clarity):
def add_two(fn):
def wrapper(self, x):
y = x + 2
return fn(self, y)
return wrapper
class Test:
#add_two
def func1(self, b):
return b
f = Test()
f.func1(5) # returns 7
This issue here isn't the decorator at all. This issue is you are using func1 and your decorator as static methods without removing the self argument. If you remove the self arguments this will work fine.
Without staticmethod decorator
class Test:
def add_two(func=None):
def wrapper_add_two(*args, **kwargs):
return func(*args, **kwargs) + 2
return wrapper_add_two
#add_two
def func1(b):
return b
print(Test.func1(10)) #12
With staticmethod decorator
Unfortunately using them in this manner stores them as unbound static methods and you need to bind them for this to work properly.
class Test:
#staticmethod
def add_two(func):
def wrapper_add_two(*args, **kwargs):
return func.__func__(*args, **kwargs) + 2
return wrapper_add_two
#add_two.__func__
#staticmethod
def func1(b):
return b
print(Test.func1(10)) #12
Running with the staticmethod decorator and without the function binding gives you
TypeError: 'staticmethod' object is not callable

Registering methods with decorator at class level

I'd like to be able to register/return methods at a class level. The closest answer I could find was here: Auto-register class methods using decorator, except it was centered on a global register and I'm looking for something specific to the class per below.
Code:
class ExampleClass:
def get_reports(self):
# return list of all method names with #report decorator
pass
def report(self):
# decorator here
pass
#report
def report_x(self):
return
#report
def report_y(self):
return
def method_z(self):
pass
where I'd like ExampleClass.get_reports() to return ['report_x', 'report_y'].
Not all reports will be preceded with report_, so there is likely no way to just look at method names. I'm trying to figure out how to do this to apply to a variety of situations, so just looking for 'report_' does not work in this context.
You can declare a Reporter class like this and use an instance as a class property. I used the __call__ override to shorten the decorator, but you can name the function report and decorate methods as #report.report
class Reporter:
def __init__(self):
# Maintain a set of reported methods
self._reported = set()
def __call__(self, fn, *args, **kwargs):
# Register method
self._reported.add(fn.__name__)
def decorate(*args, **kwargs):
return fn(*args, **kwargs)
return decorate
class ExampleClass:
report = Reporter()
def get_reports(self):
# return list of all method names with #report decorator
return list(self.report._reported)
#report
def report_x(self):
return
#report
def report_y(self):
return
def method_z(self):
pass
This turns out to be similar to Mach_Zero's answer. So key differences, this returns the methods, not the method names, and the implementation of get_reports() is somewhat simpler with the use of __iter__.
Code:
class Reports:
def __init__(self):
self.reports = []
def __call__(self, func):
self.reports.append(func)
return func
def __iter__(self):
return iter(self.reports)
class ExampleClass:
report = Reports()
#classmethod
def get_reports(cls):
# return list of all method names with #report decorator
return list(cls.report)
#report
def report_x(self):
return
#report
def report_y(self):
return
def method_z(self):
pass
Test Code:
print(ExampleClass.get_reports())
Results:
[
<function ExampleClass.report_x at 0x000000000AF7B2F0>,
<function ExampleClass.report_y at 0x000000000AF7B378>
]

how to use decorator in a class

I know there is similar question, but my scenario is somehow different: refer to codes:
class MyClass(object):
def __init__(self, log_location)
self.logs = logging(log_location) # create log object by the log_location, this object should be used by the decorator fucntion
def record_log(log_object):
""" this is the decorator function
"""
def deco(func):
def wrap(*args, **kwargs):
rs = func()
# use log object to record log
if rs:
log_object.record('success')
else:
log_object.record('fail')
return wrap
return deco
#record_log(self.logs)
def test(self):
rs = do_some_thing
if rs:
return True
return False
def main():
my_class = MyClass()
my_class.test()
But, there is an error like this:
#record_log(self.logs)
NameError: name 'self' is not defined
Hos should I use the instance attribute self.logs in a decorator function in such scenario like this??
Thanks very much!
You can not pass a reference to self or any attribute of self at this point. The #record_log line is executed (the method is decorated) before the code in main is executed, i.e. before any instance of MyClass is created at all -- in fact, even before the definition of MyClass has been completed! But remember that
#record_log(self.logs)
def test(self, n):
is actually just syntactic sugar for
test = record_log(self.logs)(test)
So one way to work around your problem would be to redefine test in your __init__, i.e.
def __init__(self, log_location)
self.logs = logging(log_location)
self.test = record_log(self.logs)(self.test)
Also note that your decorator is not passing any parameters to func and not returning the results. Also, it should probably be defined on module level (before the class).
def record_log(log_object):
def deco(func):
def wrap(*args, **kwargs):
rs = func(*args, **kwargs) # pass parameters
if rs:
log_object.record('success')
else:
log_object.record('fail')
return rs # return result
return wrap
return deco
There are several objections about your code:
deco() is redundant. You can directly return wrap from record_log().
If you only plan to decorate MyClass's methods, then there is no point in passing log_object to the decorator, as self.logs will always be used. Otherwise, consider moving the decorator to module level, as already suggested by others.
The decorated method's return value is currently lost.
The call to the decorated function does not pass self to it.
The proper code would therefore be:
class MyClass(object):
def __init__(self, log_location):
self.logs = logging(log_location)
def record_log(func):
""" this is the decorator function
"""
def wrap(self):
rs = func(self)
# use log object to record log
if rs:
print 1
self.logs.record('success')
else:
print 2
self.logs.record('fail')
return rs
return wrap
#record_log
def test(self):
rs = do_some_thing
if rs:
return True
return False

A decorator-creating class

I am writing a bunch of code that has a possibility of mutable outputs, like an arithmetic function where I could have the output be a float or an int. Basically my problem is that if I were to create a decorator for each object type I need (probably seven or eight), I would go insane with the constant repetition of:
def Int(fn):
def wrapper():
return int(fn())
return wrapper
What I want to have is a class like below that would create a decorator based on the name it's instantiated with and it would be a copy of the function above but with the appropriate type modifications.
class Decorator(object):
def __init__(self):
...
...
Int = Decorator()
# Then I can use #Int
Any help would be really appreciated. Thanks.
You cannot have Decorator know what name it will be assigned to. Assignment occurs after instantiation, so the object will have already been created by the time it is assigned a name.
You could however make a decorator that creates decorators dynamically:
from functools import wraps
def set_return_type(typeobj):
def decorator(func):
#wraps(func)
def wrapper(*args, **kwargs):
return typeobj(func(*args, **kwargs))
return wrapper
return decorator
You would then use this decorator by giving a type object argument for the type you want:
#set_return_type(int) # Causes decorated function to return ints
#set_return_type(float) # Causes decorated function to return floats
Below is a demonstration:
>>> from functools import wraps
>>> def set_return_type(typeobj):
... def decorator(func):
... #wraps(func)
... def wrapper(*args, **kwargs):
... return typeobj(func(*args, **kwargs))
... return wrapper
... return decorator
...
>>> #set_return_type(float)
... def test():
... return 1
...
>>> test()
1.0
>>>

Categories