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

Answer 3

While working on Apport, I ran into the exact same question. I experimented with multiple approaches and came up with this solution, which is the most elegant and readable one:

# Copyright (C) 2022, Benjamin Drung <bdrung@posteo.de>
#
# Permission to use, copy, modify, and/or distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

import contextlib
import typing
import unittest.mock

@contextlib.contextmanager
def wrap_object(
    target: object, attribute: str
) -> typing.Generator[unittest.mock.MagicMock, None, None]:
    """Wrap the named member on an object with a mock object.

    wrap_object() can be used as a context manager. Inside the
    body of the with statement, the attribute of the target is
    wrapped with a :class:`unittest.mock.MagicMock` object. When
    the with statement exits the patch is undone.

    The instance argument 'self' of the wrapped attribute is
    intentionally not logged in the MagicMock call. Therefore
    wrap_object() can be used to check all calls to the object,
    but not differentiate between different instances.
    """
    mock = unittest.mock.MagicMock()
    real_attribute = getattr(target, attribute)

    def mocked_attribute(self, *args, **kwargs):
        mock.__call__(*args, **kwargs)
        return real_attribute(self, *args, **kwargs)

    with unittest.mock.patch.object(target, attribute, mocked_attribute):
        yield mock

Unittest for wrap_object:

class Multiply:
    """Test class for wrap_object test cases."""

    def __init__(self, multiplier):
        self.multiplier = multiplier

    def multiply(self, x: int) -> int:
        return x * self.multiplier


class TestWrapObject(unittest.TestCase):
    def test_wrap_object_with_statement(self):
        with wrap_object(Multiply, "__init__") as mock:
            m = Multiply(7)
            self.assertEqual(m.multiply(6), 42)
        mock.assert_called_once_with(7)

Then you can use following code in your initial example:

>>> with wrap_object(Foo, "foo") as m:
...     foo = Foo()
...     print(foo.foo())
... 
foo
>>> m.assert_called_once_with()
>>> m.call_count
1
>>>