Is there a way to access the original function in a mocked method/function such that I can modify the arguments and pass it to the original functions?

Question

I'd like to modify the arguments passed to a method in a module, as opposed to replacing its return value.

I've found a way around this, but it seems like something useful and has turned into a lesson in mocking.

module.py

from third_party import ThirdPartyClass

ThirdPartyClass.do_something('foo', 'bar')
ThirdPartyClass.do_something('foo', 'baz')

tests.py

@mock.patch('module.ThirdPartyClass.do_something')
def test(do_something):
    # Instead of directly overriding its return value
    # I'd like to modify the arguments passed to this function.

    # change return value, no matter inputs
    do_something.return_value = 'foo'

    # change return value, based on inputs, but have no access to the original function
    do_something.side_effect = lambda x, y: y, x

    # how can I wrap do_something, so that I can modify its inputs and pass it back to the original function?
    # much like a decorator?

I've tried something like the following, but not only is it repetitive and ugly, it doesn't work. After some PDB introspection.. I'm wondering if it's simply due to however this third party library works, as I do see the original functions being called successfully when I drop a pdb inside the side_effect.

Either that, or some auto mocking magic I'm just not following that I'd love to learn about.

def test():
    from third_party import ThirdPartyClass
    original_do_something = ThirdPartyClass.do_something

    with mock.patch('module.ThirdPartyClass.do_something' as mocked_do_something:
        def side_effect(arg1, arg2):
            return original_do_something(arg1, 'overridden')

        mocked_do_something.side_effect = side_effect

        # execute module.py

Any guidance is appreciated!

Answer 1

You may want to use parameter wraps for the mock call. (Docs for reference.) This way the original function will be called, but it will have everything from Mock interface.

So for changing parameters called to original function you may want to try it like that:

org.py:

def func(x):
    print(x)

main.py:

from unittest import mock

import org


of = org.func
def wrapped(a):
    of('--{}--'.format(a))


with mock.patch('org.func', wraps=wrapped):
    org.func('x')
    org.func.assert_called_with('x')

result:

 --x--

Answer 2

The trick is to pass the original underlying function that you still want to access as a parameter to the function.

Eg, for race condition testing, have tempfile.mktemp return an existing pathname:

def mock_mktemp(*, orig_mktemp=tempfile.mktemp, **kwargs):
    """Ensure mktemp returns an existing pathname."""
    temp = orig_mktemp(**kwargs)
    open(temp, 'w').close()
    return temp

Above, orig_mktemp is evaluated when the function is declared, not when it is called, so all invocations will have access to the original method of tempfile.mktemp via orig_mktemp.

I used it as follows:

@unittest.mock.patch('tempfile.mktemp', side_effect=mock_mktemp)
def test_retry_on_existing_temp_path(self, mock_mktemp):
    # Simulate race condition: creation of temp path after tempfile.mktemp
    ...