Python Mocking – How to store function arguments to a mocked function in the mock that is returned?

Question:

Consider two explicit MagicMocks, such that one is used to create new mock instances by calling a method with arguments, and these mocks in turn are passed to the other mock’s method as arguments:

In [1]: from unittest.mock import MagicMock

In [2]: foo = MagicMock()

In [3]: bar = MagicMock()

In [4]: a = foo.a(1, 2)

In [5]: b = foo.b(3, 4)

In [6]: bar.func(a, b)
Out[6]: <MagicMock name='mock.func()' id='140383162348976'>

In [7]: bar.method_calls
Out[7]: [call.func(<MagicMock name='mock.a()' id='140383164249232'>, <MagicMock name='mock.b()' id='140383164248848'>)]

Note that the bar.method_calls list contains calls to the functions .a and .b, but the parameters that were passed to those functions are missing. From what I can tell, they aren’t recorded in bar at all. They can be found in foo:

In [8]: foo.method_calls
Out[8]: [call.a(1, 2), call.b(3, 4)]

But they are decoupled from their use as parameters to bar.func, thus unusable as a check that bar.func was called correctly in anything but the most trivial case (for instance, there may be many calls to foo.a unrelated to the bar.func call).

At first, I expected that the new mocks a and b would store the parameters passed, but in fact they do not, because foo.a(...) returns a new MagicMock that happens to have the name mock.a(), but call and arguments are recorded by foo. a does not store them. Same with b. Therefore when bar.func(a, b) is called, the arguments for a and b are not present, nor stored in bar.

Can the foo mock be configured in some way to create new MagicMock instances that record the parameters passed to its .a and .b methods? If not, can the code be refactored to capture the full call history in bar? Ideally lines 4 – 6 are not test code, and should remain unaware of any mocking.

EDIT: to be clear, my goal is to be able to test that the function bar.func was called with parameters foo.a(1, 2) and foo.b(3, 4). This seems to be fundamentally different from testing that the function func was called with parameters (1, 2) and (3, 4) due to the extra indirection of bar. and foo..


(The comment below was ultimately addressed in the final accepted answer, but I leave it here for posterity)

EDIT2: blhsing has offered a solution involving a subclass of MagicMock that mostly works. However there’s one case that fails:

class TraceableMock(MagicMock):
    def __call__(self, *args, **kwargs):
        child_mock = super().__call__(*args, **kwargs)
        child_mock.attach_mock(self, 'parent_mock')
        return child_mock

foo = TraceableMock()
bar = MagicMock()
a = foo.a(1, 2)
a2 = foo.b(5, 6)  # extra call to foo.a, unrelated to the upcoming bar.func() call
b = foo.b(3, 4)
bar.func(a, b)
print(bar.func.call_args.args[0].parent_mock.mock_calls)
print(bar.func.call_args.args[1].parent_mock.mock_calls)
print(bar.func.call_args.args[0].parent_mock.mock_calls == [call(1, 2)])

This outputs:

[call(1, 2), call(5, 6)]
[call(3, 4)]
False

I think this is because the Mock created for foo.a is reused, and therefore records an additional call. I can test for this:

assert call(1, 2) in bar.func.call_args.args[0].parent_mock.mock_calls

But unfortunately this does not guarantee that the call(1, 2) was actually one of the parameters to bar.func().

I could impose a condition that foo.a and foo.b are each only called once, but this is too strict, because there is no reason why these functions cannot be called multiple times, and it’s only the call to bar.func, and its parameters, that I care about in this case.

In the context of my overall problem, I’m beginning to wonder if perhaps a better is to patch in smart, custom wrapper objects, able to log their own calls, rather than trying to use Mocks.

Asked By: davidA

||

Answers:

You’re correct that as it is, a child Mock object returned by making a call to a Mock object does not have any linkage back to the Mock object that created it, thus losing the record of the call from the standpoint of the child object.

You can work around the lack of such a linkage by creating your own subclass of Mock with a wrapper __call__ method that returns a Mock object cloned from the original returning Mock object, with last the call object of the parent Mock, as well as the parent Mock object itself, stored as additional attributes:

from unittest.mock import MagicMock, call

class TraceableMock(MagicMock):
    def __call__(self, *args, **kwargs):
        child_mock = super().__call__(*args, **kwargs)
        if isinstance(child_mock, TraceableMock):
            mock = TraceableMock()
            mock.__dict__ = child_mock.__dict__.copy()
            mock.created_from = self.mock_calls[-1]
            mock.attach_mock(self, 'parent_mock')            
            return mock
        return child_mock

foo = TraceableMock()
bar = MagicMock()
a = foo.a(1, 2)
a2 = foo.a(5, 6)
b = foo.b(3, 4)
c = foo(7, 8)
bar.func(a, b)

print(bar.func.call_args.args[0].created_from)
print(bar.func.call_args.args[1].created_from)
print(c.created_from)
print(bar.func.call_args.args[0].created_from == call(1, 2))
print(bar.func.call_args.args[0].parent_mock is foo.a)

This outputs:

call(1, 2)
call(3, 4)
call(7, 8)
True
True

Demo: https://replit.com/@blhsing/CalculatingIllCustomization

Answered By: blhsing