使用unittest.mock对装饰器进行Autospec

Question

假设我有一个简单的装饰器方法，像下面这样：

def my_decorator(fn):
  def _wrapper(*args, **kwargs):
    print 'Calling decorated function'
    fn(*args, **kwargs)
  return _wrapper

class Foo(object):
  @my_decorator
  def incr(self, x):
    return x+1

这个装饰器会“抹去”方法的签名，以便于自动规范检查：

>>> mock_foo = mock.create_autospec(Foo, instance=True)
>>> mock_foo.incr(1, 2, 3, 4)
<MagicMock name='mock.incr()' id='23032592'>

这应该会引发错误：

TypeError: <lambda>() takes exactly 2 arguments (5 given)

我曾经因为关键字参数的拼写错误而出现过这样的bug。

有没有办法编写这个装饰器（或者给自动规范检查一个“提示”），让这些错误能够被捕捉到呢？

Answer 1

一个同事给我推荐了这个装饰器库

from decorator import decorator

@decorator
def my_decorator(fn, *args, **kwargs):
  print 'Calling decorated function'
  return fn(*args, **kwargs)

class Foo(object):
  @my_decorator
  def incr(self, x):
    return x+1

@decorator这个东西能做到一些很厉害的事情，它可以让你在使用@my_decorator的时候，不会让被装饰的函数的签名被隐藏。

Answer 2

你可以使用functools.wraps，mock.create_autospec也能理解它：

from functools import wraps
from unittest import mock


def my_decorator(fn):
   @wraps(fn)
   def _wrapper(*args, **kwargs):
       print('Calling decorated function')
       fn(*args, **kwargs)
   return _wrapper


class Foo(object):
    @my_decorator
    def incr(self, x):
        return x + 1


if __name__ == '__main__':
    mock_foo = mock.create_autospec(Foo, instance=True)
    mock_foo.incr(1, 2, 3, 4)

如果你把上面的代码放到一个文件里运行，你会看到最后一行的错误信息：

TypeError: 参数太多了

如果没有使用@wraps，这个脚本会正常结束，返回代码是0。

Answer 3

我觉得autospec不能直接做到这一点。不过，你可以在装饰器里做一些小技巧，让你能测试未被装饰的函数。如果你让你的装饰器保存一个未装饰函数的引用：

def my_decorator(fn):
  def _wrapper(*args, **kwargs):
    print 'Calling decorated function'
    fn(*args, **kwargs)
  _wrapper._orig = fn
  return _wrapper

你可以通过被模拟的装饰函数来访问它：

>>> mock_incr = mock.create_autospec(Foo.incr)
>>> mock_incr(1,3,4,5,5)               # Decorated function doesn't fail.
<MagicMock name='mock()' id='8734864'>
>>> mock_incr._orig(1,3,4,5,5)         # But the original does, which is what we want
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/lib64/python2.6/site-packages/mock.py", line 954, in __call__
    _mock_self._mock_check_sig(*args, **kwargs)
TypeError: <lambda>() takes exactly 3 arguments (6 given)
>>> mock_incr._orig(1,3)
<MagicMock name='mock._orig()' id='8739664'>

不过，如果你对整个实例使用autospec，这样就不行了。我也不太清楚为什么。

>>> mock_foo = mock.create_autospec(Foo, instance=True)
>>> mock_foo.incr(1,3,4,5)             # We expect this to not raise an exception
<MagicMock name='mock.incr2()' id='8758416'>
>>> mock_foo.incr._orig(1,3,4,5)       # But we were hoping this would :(
<MagicMock name='mock.incr._orig()' id='8740624'>

另外，值得一提的是Venusian，它可以改变装饰器与被装饰方法绑定的方式，专门为了解决这个问题。不过，这可能比你想要的要复杂一些。