Python mock: wrap instance method

2020-07-23 06:21发布

问题:

What I would like: Ensure that all instances of Foo that are created inside the with statement have their foo instance method wrapped in a MagicMock via wraps=Foo.foo. The reason I want this is so that I can track call_count on the method foo for all instances of Foo that are created. Now that I say it like that it seems kind of impossible...

>>> from mock import patch
...
... class Foo(object):
...
...     def foo(self):
...         return "foo"
...
... with patch("__main__.Foo.foo", wraps=Foo.foo) as m:
...     foo = Foo()
...     print(foo.foo())

Traceback (most recent call last):
  File "a.py", line 12, in <module>
    print(foo.foo())
  File "/disk/software/lib/python27/mock/mock.py", line 1062, in __call__
    return _mock_self._mock_call(*args, **kwargs)
  File "/disk/software/lib/python27/mock/mock.py", line 1132, in _mock_call
    return self._mock_wraps(*args, **kwargs)
TypeError: unbound method foo() must be called with Foo instance as first argument (got nothing instead)

The problem The mocked foo method isn't bound to the foo instance created via foo = Foo() because it's wrapping the unbound method Foo.foo. Does anyone know how to ensure that the mocked method is bound to an instance?

What I already know:

>>> foo = Foo()
... with patch.object(foo, "foo", wraps=foo.foo) as m:
...     print(foo.foo())
"foo"

But this doesn't satisfy my constraint that the object must be instantiated inside the patch context.

回答1:

The problem with my proposed and incorrect solution above

with patch("__main__.Foo.foo", wraps=Foo.foo) as m:
    ...

is that the foo method on Foo is mocked such that it wraps the unbound method Foo.foo, which naturally doesn't work, because the unbound method Foo.foo has no idea which instance it's attached to when called later on.

The simplest solution I could think of

from mock import patch, MagicMock

class Foo:

    def foo(self):
        return "foo"

class MockFoo(Foo):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Every instance of MockFoo will now have its `foo` method 
        # wrapped in a MagicMock
        self.foo = MagicMock(wraps=self.foo)

with patch("__main__.Foo", MockFoo) as m:
    foo = Foo()
    print(foo.foo())
    assert foo.foo.call_count == 1