CODEWITHSUNDEEP
pythonunit-testingmocking
Daniel Kats
Daniel Kats

Reputation: 5554

Python3 mock replace function with another function

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?

Upvotes: 22

Views: 33013

Answers (2)

chepner
chepner

Reputation: 531125

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.

Upvotes: 32

wwii
wwii

Reputation: 23753

AttributeError: 'function' object has no attribute 'assert_called_once

Sounds like you are trying to path with a regular function. Patch with a MagicMock object whose return value is the result of calling fake_bar.

from unittest.mock import patch, MagicMock

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

def bar(arg):
    return 2

def fake_bar():
    #calculate 2!
    return 2

def test_foo():
    with patch("__main__.bar", MagicMock(return_value=fake_bar())) as mock_bar:
        arg1,arg2 = 'a','b'
        foo(arg1, arg2)
        mock_bar.assert_called_once_with(arg1)
##        mock_bar.assert_called_once_with('w')
        print(mock_bar.call_args)

Variants:

def test_foo():
    with patch("__main__.bar", autospec=True) as mock_bar:
        arg1,arg2 = 'a','b'
        mock_bar.return_value = fake_bar()
        foo(arg1, arg2)
        mock_bar.assert_called_once_with(arg1)
        print(mock_bar.call_args)

Upvotes: 1

Related Questions