Python 3: unittest.mock how to specify different return values for specific inputs?

Question

Say I have a mock object (let it be Mock or MagicMock). I want to mock one of its method to return one value for a specific input and to return another value for another specific input. How do I do it without having to care about the order the input is sent to the method?

Pseudo code:

def test_blah():
    o = MagicMock()

    # How to do I these?
    o.foo.return_value.on(1) = 10
    o.foo.return_value.on(2) = 20

    assert o.foo(2) == 20
    assert o.foo(1) == 10

Answer 1

Since you need to dynamically change the return values, you need to use side_effect and not the return_value. You may map the desired return values to the expected inputs and use the .get function of this mapping/dictionary as a side effect:

return_values = {1: 10, 2: 20}
o.foo.side_effect = return_values.get

Demo:

In [1]: from unittest import mock

In [2]: m = mock.MagicMock()

In [3]: return_values = {1: 10, 2: 20}

In [4]: m.foo.side_effect = return_values.get

In [5]: m.foo(1)
Out[5]: 10

In [6]: m.foo(2)
Out[6]: 20