Python mock return default value after initial predefined values returned

Question

I'm mocking a message queue, where I want it to return 2 messages and then return None each time it is called (to simulate that the message queue is now empty). However, after the 3rd message call I get an error from mocking framework saying that there is no "next" value.

Here's an example:

from mock import MagicMock

mock_message_queue = MagicMock(name='message_queue')
expected_messages = ['message 1', 'message 2', None]
mock_message_queue.fetch.side_effect = expected_messages

print('mock_message_queue.fetch() = {}'.format(mock_message_queue.fetch()))
print('mock_message_queue.fetch() = {}'.format(mock_message_queue.fetch()))
print('mock_message_queue.fetch() = {}'.format(mock_message_queue.fetch()))
print('mock_message_queue.fetch() = {}'.format(mock_message_queue.fetch()))

Output:

mock_message_queue.fetch() = message 1
mock_message_queue.fetch() = message 2
mock_message_queue.fetch() = None
Traceback (most recent call last):
  File "/home/wsi/.PyCharm2017.1/config/scratches/mock_message_bus.py", line 10, in <module>
    print('mock_message_queue.fetch() = {}'.format(mock_message_queue.fetch()))
  File "/usr/local/lib/python3.5/dist-packages/mock/mock.py", line 1062, in __call__
    return _mock_self._mock_call(*args, **kwargs)
  File "/usr/local/lib/python3.5/dist-packages/mock/mock.py", line 1121, in _mock_call
    result = next(effect)
StopIteration

Process finished with exit code 1

So basically I want it to always return None after it returns 'message 2'. Anyone know how to do this?

score 2 · Accepted Answer · answered Jun 21 '17 at 12:49

The issue is that you need the last element of the list of values to repeat endlessly. You can use itertools.repeat for this, along with itertools.chain to join it onto the list of actual values.

expected_messages = itertools.chain(['message 1', 'message 2'], itertools.repeat(None))

Python mock return default value after initial predefined values returned

1 Answers1