使用unittest.mock通过装饰器自动指定

Question

Suppose I have a simple decorated method like the following:

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

The decorator "erases" the method signature for autospec'ing purposes:

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

This should raise:

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

I've had bugs like this creep through due to typos in keyword arguments.

Is there any way to write the decorator (or give a "hint" to autospec) so that these kinds of errors will be caught?

Answer 1

I don't think autospec can do this directly. You can do a little hackery in the decorator to make it possible to test your undecorated function, though. If you make your decorator save a reference to the undecorated function:

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

You can access it via the mocked decorated function:

>>> 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'>

This doesn't work if you autospec the whole instance, though. Not sure why.

>>> 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'>

Also worth noting is Venusian , which can change the way decorators get bound to decorated methods specifically to address this use-case. Might be more heavyweight than you want, though.

Answer 2

A coworker pointed me to the decorator library

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 implements the right magic to not hide the wrapped function's signature behind @my_decorator.

Answer 3

You can use functools.wraps, mock.create_autospec understands it:

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)

if you place code above into file and run, you'll see trace with last line:

TypeError: too many positional arguments

Without @wraps script ends with exit code 0.