While implementing unit test in python, we need to replace some of the functions with the mock objects and perform the necessary assertions. These mock objects can be configured to specify the return values. In addition to this, we can also specify the return values based on the arguments of the function that has been mocked.

For instance, let's assume ee have a function below that needs to be tested

def dummy_function(param1, param2):
    return param1 + param2

Normally, the way we mock this function is

test_value = MagicMock(return_value="tested!")

Or if we want to call a function multiple times and return different output in each call it can be achieve as

test_value = MagicMock(side_effect=["result1", "result2"])

This will just return the value of the function without caring about the parameters provided to it.

However, we might also want to test based on the arguments of a functions such that when function args is 'foo' then function return 'bar' else it returns something else.

In such situation, we can take the advantage of side_effect. Using side_effect we can control the behavior of our mock objects and instruct it to return something when the value is something. Let's look at how we can do this.

One way is to create a function with your own side_effect.

from unittest.mock import MagicMock, patch

def args_based_return(*args, **kwargs):
    if args == (10, 10):
        return 20
    elif args  == (20, 20):
        return 40
    else:
        return Exception("exception occurred")

m = MagicMock(side_effect=args_based_return)
m(10, 10)
>> 20

m(40, 40)
>> exception occurred

Another way is by using the lambda function

m = MagicMock()
m.side_effect = lambda *x: {
        (10, 10): 20,
        (20, 20): 40,
        (30, 30):  60
    }[x]

print(m(30, 30))
>> 60

print(m(20, 20))
>> 40