117 lines
3.1 KiB
Python
117 lines
3.1 KiB
Python
from __future__ import print_function
|
|
|
|
import unittest
|
|
|
|
import wrapt
|
|
|
|
class TestDecorator(unittest.TestCase):
|
|
|
|
def test_no_parameters(self):
|
|
_args = (1, 2)
|
|
_kwargs = {'one': 1, 'two': 2}
|
|
|
|
@wrapt.decorator
|
|
def _decorator(wrapped, instance, args, kwargs):
|
|
return wrapped(*args, **kwargs)
|
|
|
|
@_decorator
|
|
def _function(*args, **kwargs):
|
|
return args, kwargs
|
|
|
|
result = _function(*_args, **_kwargs)
|
|
|
|
self.assertEqual(result, (_args, _kwargs))
|
|
|
|
def test_instance_method_as_decorator(self):
|
|
_args = (1, 2)
|
|
_kwargs = {'one': 1, 'two': 2}
|
|
|
|
class Instance(object):
|
|
def __init__(self):
|
|
self.count = 0
|
|
@wrapt.decorator
|
|
def decorator(self, wrapped, instance, args, kwargs):
|
|
self.count += 1
|
|
return wrapped(*args, **kwargs)
|
|
|
|
instance = Instance()
|
|
|
|
@instance.decorator
|
|
def _function(*args, **kwargs):
|
|
return args, kwargs
|
|
|
|
result = _function(*_args, **_kwargs)
|
|
|
|
self.assertEqual(result, (_args, _kwargs))
|
|
self.assertEqual(instance.count, 1)
|
|
|
|
result = _function(*_args, **_kwargs)
|
|
|
|
self.assertEqual(result, (_args, _kwargs))
|
|
self.assertEqual(instance.count, 2)
|
|
|
|
def test_class_method_as_decorator(self):
|
|
_args = (1, 2)
|
|
_kwargs = {'one': 1, 'two': 2}
|
|
|
|
class Instance(object):
|
|
count = 0
|
|
@wrapt.decorator
|
|
@classmethod
|
|
def decorator(cls, wrapped, instance, args, kwargs):
|
|
cls.count += 1
|
|
return wrapped(*args, **kwargs)
|
|
|
|
@Instance.decorator
|
|
def _function(*args, **kwargs):
|
|
return args, kwargs
|
|
|
|
result = _function(*_args, **_kwargs)
|
|
|
|
self.assertEqual(result, (_args, _kwargs))
|
|
self.assertEqual(Instance.count, 1)
|
|
|
|
result = _function(*_args, **_kwargs)
|
|
|
|
self.assertEqual(result, (_args, _kwargs))
|
|
self.assertEqual(Instance.count, 2)
|
|
|
|
def test_class_type_as_decorator(self):
|
|
_args = (1, 2)
|
|
_kwargs = {'one': 1, 'two': 2}
|
|
|
|
@wrapt.decorator
|
|
class ClassDecorator(object):
|
|
def __call__(self, wrapped, instance, args, kwargs):
|
|
return wrapped(*args, **kwargs)
|
|
|
|
@ClassDecorator
|
|
def _function(*args, **kwargs):
|
|
return args, kwargs
|
|
|
|
result = _function(*_args, **_kwargs)
|
|
|
|
self.assertEqual(result, (_args, _kwargs))
|
|
|
|
def test_class_type_as_decorator_args(self):
|
|
_args = (1, 2)
|
|
_kwargs = {'one': 1, 'two': 2}
|
|
|
|
@wrapt.decorator
|
|
class ClassDecorator(object):
|
|
def __init__(self, arg):
|
|
assert arg == 1
|
|
def __call__(self, wrapped, instance, args, kwargs):
|
|
return wrapped(*args, **kwargs)
|
|
|
|
@ClassDecorator(arg=1)
|
|
def _function(*args, **kwargs):
|
|
return args, kwargs
|
|
|
|
result = _function(*_args, **_kwargs)
|
|
|
|
self.assertEqual(result, (_args, _kwargs))
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|