Как заставить mock-объект в Python возвращать разные значения при последовательных вызовах

В стандартной библиотеке unittest.mock для этой цели используется атрибут side_effect. Он позволяет гибко управлять поведением mock-объекта при каждом вызове.

1. Использование итерируемого объекта (например, списка)

Если передать в side_effect список, кортеж или любой другой итератор, при каждом вызове mock-объекта будет возвращаться следующий элемент этого итератора.

from unittest.mock import Mock

mock_obj = Mock()
mock_obj.side_effect = [10, 'hello', True, ValueError('An error')]

# Последовательные вызовы
print(mock_obj())  # 10
print(mock_obj())  # 'hello'
print(mock_obj())  # True

# Четвертый вызов вызовет исключение, указанное в списке
try:
    mock_obj()
except ValueError as e:
    print(e) # An error

# Пятый вызов вызовет StopIteration, так как элементы закончились
try:
    mock_obj()
except StopIteration:
    print('Iterator is exhausted')

2. Использование функции (callable)

Передача функции в side_effect позволяет реализовать динамическую логику. Эта функция будет вызываться с теми же аргументами, что и mock-объект.

from unittest.mock import Mock

def dynamic_return_value(*args, **kwargs):
    # Возвращаем значение в зависимости от переданных аргументов
    if 'id' in kwargs and kwargs['id'] == 1:
        return {'status': 'success'}
    return {'status': 'not_found'}

api_mock = Mock()
api_mock.side_effect = dynamic_return_value

print(api_mock(id=1))  # {'status': 'success'}
print(api_mock(id=2))  # {'status': 'not_found'}
print(api_mock())      # {'status': 'not_found'}

Краткий итог:

• Список: для предопределенной последовательности возвращаемых значений или исключений.
• Функция: для сложной логики, когда возвращаемое значение зависит от аргументов вызова.