返回初始预定义值后,Python模拟返回默认值

时间:2017-06-21 12:27:35

标签: python unit-testing mocking python-mock

我正在模拟一个消息队列,我希望它返回2条消息,然后在每次调用时返回None(模拟消息队列现在为空)。但是,在第三次消息调用后,我从模拟框架中得到一个错误,说没有" next"值。

以下是一个例子:

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()))

输出:

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

所以基本上我希望它在返回&#39;消息2&#39;后总是返回None。有谁知道怎么做?

1 个答案:

答案 0 :(得分:2)

问题是您需要值列表的最后一个元素无休止地重复。您可以使用itertools.repeat以及itertools.chain将其加入实际值列表中。

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