Python 3:unittest.mock如何为特定的输入指定不同的return值?

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

假设我有一个模拟对象(让它是 Mock 或 MagicMock)。我想将其方法之一模拟为特定输入的 return 一个值和另一个特定输入的 return 另一个值。如何做到这一点而不必关心输入发送到方法的顺序?

伪代码:

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

由于您需要动态更改 return 值,因此您需要使用 side_effect 而不是 return_value。您可以将所需的 return 值映射到预期的输入,并使用此 mapping/dictionary 的 .get 函数作为副作用:

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

演示:

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