Python mock: wrap instance method

Question

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.

Answer 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

Answer 2

This is so nasty with Python mocks so that I ended up building a custom patch implementation (extend with other features if required).

import contextlib

class Patcher:
    UNCHANGED_RET = object()

    def __init__(self):
        self.call_count = 0
        self.return_value = Patcher.UNCHANGED_RET


@contextlib.contextmanager
def patch(klass, method_name):
    patcher = Patcher()
    orig_method = getattr(klass, method_name)

    def new_method(myself, *args, **kwargs):
        patcher.call_count += 1
        orig_return_value = orig_method(myself, *args, **kwargs)

        if patcher.return_value != Patcher.UNCHANGED_RET:
            return patcher.return_value

        return orig_return_value

    setattr(klass, method_name, new_method)
    yield patcher
    setattr(klass, method_name, orig_method)

Use as follows:

class MyClass:
    def f(self):
        return 42


x = MyClass()
with patch(MyClass, 'f') as f_patcher:
    y = MyClass()  # inside or outside -- does not matter
    assert x.f() == 42
    assert f_patcher.call_count == 1
    f_patcher.return_value = 7
    assert y.f() == 7
    assert f_patcher.call_count == 2