Skip to main content

How to Mock Multiple Return Values in Python Unit Tests

When unit testing, you often need to mock functions or methods to return different values on successive calls.

This guide explains how to mock multiple return values using the unittest.mock module in Python, covering the side_effect attribute of mock objects and the patch() decorator, as well as ways to handle the end of a sequence of return values.

Using side_effect with a List​

The most direct way to mock multiple return values is to set the side_effect attribute of a Mock object to a list:

from unittest.mock import Mock

m = Mock()
m.side_effect = ['tutorial', 'reference', 'com']

print(m())  # Output: tutorial
print(m())  # Output: reference
print(m())  # Output: com
  • Each time the mock object (m in this case) is called, it returns the next value from the side_effect list.

Handling StopIteration​

Once the side_effect list is exhausted, calling the mock again will raise a StopIteration exception:

from unittest.mock import Mock

m = Mock()
m.side_effect = ['tutorial', 'reference', 'com']

print(m())  # Output: tutorial
print(m())  # Output: reference
print(m())  # Output: com

try:
    print(m())  # Raises StopIteration
except StopIteration:
    print("StopIteration raised!")
  • You can catch this exception using try/except, or you may want to provide a default after the list of return values are exhausted.

Using itertools.chain for default after exhaustion​

    from unittest.mock import Mock
import itertools

m = Mock()

m.side_effect = itertools.chain(
    ['tutorial', 'reference'],
    itertools.repeat('com') # Repeats "com" forever
)

print(m())  # Output: tutorial
print(m())  # Output: reference
print(m())  # Output: com
print(m())  # Output: com
  • The itertools.chain() method can be used to chain iterators.
  • The itertools.repeat() method will create an iterator that returns a constant value.

Using side_effect with a List (upon instantiation)​

You can also define the return values during the creation of the Mock object.

from unittest.mock import Mock
m = Mock(side_effect=['tutorial', 'reference', 'com'])
print(m())  # Output: tutorial
print(m())  # Output: reference
print(m())  # Output: com

Using side_effect with MagicMock​

If your mocked object needs to behave like a real object (with dunder methods defined), then you can use the MagicMock class instead.

from unittest.mock import MagicMock

m = MagicMock(side_effect=['tutorial', 'reference', 'com'])
print(m())  # Output: tutorial
print(m())  # Output: reference
print(m())  # Output: com
  • MagicMock is a subclass of Mock that implements default behaviors for magic methods.

Using patch() with side_effect​

The patch() decorator (or context manager) is commonly used to replace a function or method with a mock during a test. You can combine patch() with side_effect to control the return values:

example.py
def greet(name):
    return f'hello {name}'
test_example.py
from unittest.mock import patch
import example  # Import the module containing the function

@patch('example.greet', side_effect=['hello Alice', 'hello Anna', 'hello Carl'])
def test_greet_function(mock_greet): # The mock is passed as an argument
    assert example.greet('Alice') == 'hello Alice'
    mock_greet.assert_called_with('Alice')

    assert example.greet('Anna') == 'hello Anna'
    mock_greet.assert_called_with('Anna')
    assert example.greet('Carl') == 'hello Carl'
    mock_greet.assert_called_with('Carl')

    print(mock_greet)
    print(mock_greet.called)
    assert mock_greet.call_count == 3


test_greet_function()
  • @patch('example.greet', ...): This decorator replaces the greet function within the example module with a Mock object. The important thing is that you're patching where it's used, not where it's defined.
  • side_effect=['hello Alice', ...]: This sets the side_effect of the mock to the provided list.
  • def test_greet_function(mock_greet):: The mock object is passed as an argument (mock_greet in this case) to the test function.
  • example.greet(...): Inside the test, you call the original function as if it were not mocked. The mocking happens behind the scenes because of the patch decorator.
  • You can use mock_greet.assert_called_with() to make assertions about how the mocked function was called.

Table of Contents