Python mocking: how do I change behavior based on input values?

Question

I have a class that I am trying to cover with tests in python.

I have a dependency that's being injected into the class.

class UnderTest:
    def __init__(self, dependency):
        self.dependency = dependency

I don't really care about the internals of the dependency here, and want to mock it.

So I instantiate the class in my tests, injecting the dependency:

dependency = MagicMock()
dependency.some_func = MagicMock(return_value='blue')
under_test = UnderTest(dependency)

Later when I want to test the class UnderTest, I want the dependency to return a different value based on what parameters were passed to it. So in the code under test, I might have something like

value = dependency.some_func('a')

but I also want the dependency to return something else when called with a different value.

value = dependency.some_func('b')

Ideally (and I have seen this in other frameworks), I would be able to configure the mock to return different values, for example (how I would like it to work)

dependency.some_func = MagicMock([
    {'called_with': 'a', 'return_value': 'blue'}, 
    {'called_with': 'b', 'return_value': 'green'}
])

Answer 1

You can use the side_effect attribute to set a callable that returns different values according to its arguments.

results = {'a': 'blue', 'b': 'green'}
dependency.some_func = MagicMock(side_effect=lambda arg: results.get(arg, DEFAULT))

(The DEFAULT singleton is used to signal that the normal mock return value is used for arguments not in the dict.)