Mocking python function based on input arguments

Question

We have been using Mock for python for a while.

Now, we have a situation in which we want to mock a function

def foo(self, my_param):
    #do someth         


        

Answer 1

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`

Answer 2

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

Answer 3

If side_effect is a function then whatever that function returns is what calls to the mock return. The side_effect 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(value):
...     return value + 1
...
>>> m = MagicMock(side_effect=side_effect)
>>> m(1)
2
>>> m(2)
3
>>> m.mock_calls
[call(1), call(2)]

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

Answer 4

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!

Answer 5

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

Answer 6

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.