Mocking python function based on input arguments

If side_effect_func is a function then whatever that function returns is what calls to the mock return. The side_effect_func function is called with the same arguments as the mock. This allows you to vary the return value of the call dynamically, based on the input:

>>> def side_effect_func(value):
... return value + 1
...
>>> m = MagicMock(side_effect=side_effect_func)
>>> m(1)
2
>>> m(2)
3
>>> m.mock_calls
[call(1), call(2)]

http://www.voidspace.org.uk/python/mock/mock.html#calling

As indicated at Python Mock object with method called multiple times

A solution is to write my own side_effect

def my_side_effect(*args, **kwargs):
 if args[0] == 42:
 return "Called with 42"
 elif args[0] == 43:
 return "Called with 43"
 elif kwargs['foo'] == 7:
 return "Foo is seven"
mockobj.mockmethod.side_effect = my_side_effect

That does the trick

Side effect takes a function (which can also be a lambda function), so for simple cases you may use:

m = MagicMock(side_effect=(lambda x: x+1))

If you "want to return a fixed value when the input parameter has a particular value", maybe you don't even need a mock and could use a dict along with its get method:

foo = {'input1': 'value1', 'input2': 'value2'}.get
foo('input1') # value1
foo('input2') # value2

This works well when your fake's output is a mapping of input. When it's a function of input I'd suggest using side_effect as per Amber's answer.

You can also use a combination of both if you want to preserve Mock's capabilities (assert_called_once, call_count etc):

self.mock.side_effect = {'input1': 'value1', 'input2': 'value2'}.get

I've ended up here looking for "how to mock a function based on input arguments" and I finally solved this creating a simple aux function:

def mock_responses(responses, default_response=None):
 return lambda input: responses[input] if input in responses else default_response

Now:

my_mock.foo.side_effect = mock_responses(
 {
 'x': 42, 
 'y': [1,2,3]
 })
my_mock.goo.side_effect = mock_responses(
 {
 'hello': 'world'
 }, 
 default_response='hi')
...
my_mock.foo('x') # => 42
my_mock.foo('y') # => [1,2,3]
my_mock.foo('unknown') # => None
my_mock.goo('hello') # => 'world'
my_mock.goo('ey') # => 'hi'

Hope this will help someone!

Although side_effect can achieve the goal, it is not so convenient to setup side_effect function for each test case.

I write a lightweight Mock (which is called NextMock) to enhance the built-in mock to address this problem, here is a simple example:

from nextmock import Mock
m = Mock()
m.with_args(1, 2, 3).returns(123)
assert m(1, 2, 3) == 123
assert m(3, 2, 1) != 123

It also supports argument matcher:

from nextmock import Arg, Mock
m = Mock()
m.with_args(1, 2, Arg.Any).returns(123)
assert m(1, 2, 1) == 123
assert m(1, 2, "123") == 123

Hope this package could make testing more pleasant. Feel free to give any feedback.

You can also use partial from functools if you want to use a function that takes parameters but the function you are mocking does not. E.g. like this:

def mock_year(year):
 return datetime.datetime(year, 11, 28, tzinfo=timezone.utc)

@patch('django.utils.timezone.now', side_effect=partial(mock_year, year=2020))

This will return a callable that doesn't accept parameters (like Django's timezone.now()), but my mock_year function does.

Just to show another way of doing it:

def mock_isdir(path):
 return path in ['/var/log', '/var/log/apache2', '/var/log/tomcat']
with mock.patch('os.path.isdir') as os_path_isdir:
 os_path_isdir.side_effect = mock_isdir

You can also use @mock.patch.object:

Let's say a module my_module.py uses pandas to read from a database and we would like to test this module by mocking pd.read_sql_table method (which takes table_name as argument).

What you can do is to create (inside your test) a db_mock method that returns different objects depending on the argument provided:

def db_mock(**kwargs):
 if kwargs['table_name'] == 'table_1':
 # return some DataFrame
 elif kwargs['table_name'] == 'table_2':
 # return some other DataFrame

In your test function you then do:

import my_module as my_module_imported
@mock.patch.object(my_module_imported.pd, "read_sql_table", new_callable=lambda: db_mock)
def test_my_module(mock_read_sql_table):
 # You can now test any methods from `my_module`, e.g. `foo` and any call this 
 # method does to `read_sql_table` will be mocked by `db_mock`, e.g.
 ret = my_module_imported.foo(table_name='table_1')
 # `ret` is some DataFrame returned by `db_mock`

• python
• unit-testing
• mocking
• mockito