Python3 mock replace function with another function

Question

How do I mock a function with another function in python, that will also provide a mock object?

I have code that does something like:

def foo(arg1, arg2):
    r = bar(arg1)
    # does interesting things

I would like to substitute an implementation of the bar function, have it return sane values, and make sure it's called with the right arguments.

I've tried this, where fake_bar is my simple replacement for bar:

from unittest.mock import patch

def test_foo():
   with patch("x.bar", fake_bar) as mock_bar:
      actual = foo(arg1, arg2)
      assert actual == expected
      mock_bar.assert_called_once_with(arg1)

However, I get this error:

AttributeError: 'function' object has no attribute 'assert_called_once'

I found this question which suggests using create_autospec. However, mock.create_autospec does not preserve return values.

The python documentation suggests doing something like:

mock_function = create_autospec(function, return_value='fishy')

However, my faked version of bar has computed return values that are not static and not so simple to put into an inline lambda.

I feel like I'm missing something obvious. Is there a way to easily both mock a function and replace its implementation with something more complex than a static return value?

Answer 1

Use the wraps keyword argument:

def test_foo():
   with patch("x.bar", wraps=fake_bar) as mock_bar:
      actual = foo(arg1, arg2)
      assert actual == expected
      mock_bar.assert_called_once_with(arg1)

Here, x.bar is the fully qualified name of the function you are replacing, for example 'bigbusiness.account.AccountRepository.insert'.

mock_bar will still be a Mock (as is the default), but that Mock will wrap the function fake_bar, so the return value of mock_bar is the return value of fake_bar, rather than another Mock instance.