Unittest.mock.patch - record result of wrapped function call

[mimre25]

Feature or enhancement

Proposal:

I would like to access the return value of a function call wrapped with unittest.mock.patch, that I don't call myself directly.

Motivation

Let's say I have a framework that supports some "before" hooks to allow users to customize some data flowing through the framework.
One could write hook functions for such a framework and publish them.
When doing so, it would make integration tests easier if one could access the return value of a function wrapped with a MagicMock.

# framework.py
def init(before_hook):
    ... # register the update hook that is later used before doing some action to allow customization of the action

# my_plugin.py
def my_before_hook(some_data: dict) -> dict:
    ... #do something with some_data eg adding a key
    some_data["foo"] = "bar"
    return some_data # return the altered data


# test_my_plugin.py
import my_plugin
import framework

def test_my_before_hook():
    with unittest.mock.patch("my_plugin.my_before_hook",  wraps=my_plugin.my_before_hook) as m:
        framework.init(before_hook=m)
        ... #do something that runs the framework
        assert m.called
        assert ... # assert m's return value has the key "foo" with the value "bar"

Unfortunately, currently the second assert statement is not really possible - at least I couldn't find anything in the docs, debugger, and LLMs just hallucinated on this one 😅.

Alternatives

There are two alternatives I have considered:

1. Digging into the framework and finding a mock target - while this would work, it can be quite tedious, especially when the framework is big
2. Wrapping the function in question manually - Definitely possible, but seems to introduce an extra step every time this is necessary and moves the assertion to a an odd place:

# test_my_plugin.py
import my_plugin
import framework

def foo(*args, **kwargs)
    result = my_plug.my_before_hook(*args, **kwargs)
    assert result["foo"] == "bar" # this only works if the framework doesn't handle AssertionErrors
    return result


def test_my_before_hook():
    with unittest.mock.patch("my_plugin.my_before_hook",  wraps=foo) as m:
        framework.init(before_hook=m)
        ... #do something that runs the framework
        assert m.called

Proposed Solution

Record the return value of a "wraps" call and make it available on the call tuple.

I know this could be somewhat problematic due to backwards compatibility, but I'm not entirely sure.

---

I dove a little bit into the source code and don't think it will be too big of a change both in terms of code and functionality - please correct me if I'm wrong.

I'd be happy to contribute if this is enhancement is desired 🙂

Has this already been discussed elsewhere?

This is a minor feature, which does not need previous discussion elsewhere

Links to previous discussion of this feature:

No response

[tirkarthi]

I am not sure wraps should record calls. Perhaps this might be achieved with a side_effect ? I also see you are patching my_plugin.my_before_hook and wrapping it again with my_plugin.my_before_hook which feels counter intuitive.

from unittest.mock import patch, MagicMock


def init():
    my_func()


def wrapper():
    some_data = {"foo": "bar"}
    return some_data


def my_func():
    some_data = {"foo": "my_func"}
    return some_data


with patch("__main__.my_func", side_effect=wrapper) as m:
    init()
    assert m.called
    assert my_func() == {"foo": "bar"}
    assert m() == {"foo": "bar"}

[mimre25]

In your code snippet, how would you ensure that the return value of my_func is the expected one within init?

I also see you are patching my_plugin.my_before_hook and wrapping it again with my_plugin.my_before_hook which feels counter intuitive.

This was done for the lack of "alternative" to show the problem.

Below you can see a working implementation:

def test_my_before_hook():
    result = []

    def foo(*args, **kwargs):
        res = my_plugin.my_before_hook(*args, **kwargs)
        results.append(res)
        return res

    with unittest.mock.patch("my_plugin.my_before_hook",  wraps=foo) as m:
        framework.init(before_hook=m)
        ... #do something that runs the framework
        assert m.called
        assert results[0]["foo"] == "bar" # assert that the return value is as expected