`unittest.mock`: Track return values of calls?

I’m using the unittest library, and I find it nice that I can retrieve the arguments of all the calls to a method through the call_args_list attribute.

Now I wonder if it’s possible to support retrieving the return values of all the calls, too.

Here’s the background:

I’m testing a function foo which takes a bar object as input.
The bar object has a method process, which takes a variable state as input, and returns an updated state variable.
The foo function calls bar.process several times, each time feeding the last time’s output as input.

I’d like to write a unit test that verifies the input to bar.process on each call is equal to the last call’s output.
Because the foo function relies on the return value of bar.process to work correctly, I cannot mock it to return predefined values.
Instead, I’m creating a Mock object that wraps a real bar object, so it retains the bar object’s behavior:

mock_bar = Mock(wraps=bar)
foo(mock_bar)

After this, I can retrieve the input arguments to each call of mock_bar.process as:

inputs = [args_list.args[0] for args_list in mock_bar.process.call_args_list]

If I could do something like

outputs = mock_bar.process.call_return_value_list

I could then assert that inputs[1:] == outputs[:-1].

From the documents, it appears that mocks only track the input arguments to method calls, but not the return values.
Did I miss something? Or could this be supported in the future?

Although it would be a nice feature to have indeed, currently you can achieve this by subclassing Mock and overriding __call__ to keep track of the return values yourself:

class ReturnTrackingMock(Mock):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.call_return_value_list = []

    def __call__(self, *args, **kwargs):
        value = super().__call__(*args, **kwargs)
        self.call_return_value_list.append(value)
        return value

so that:

_count = 0
def count():
    global _count
    _count += 1
    return _count

mock_count = ReturnTrackingMock(wraps=count)
for _ in range(3):
    mock_count()
print(mock_count.call_return_value_list)

outputs:

[1, 2, 3]
1 Like

Great idea! Thanks for the reply :slight_smile: