模擬 side_effect 迭代器在耗盡后可以重置嗎？

Question

mock.reset_mock()不會重置副作用迭代器。 有沒有辦法在不再次創建模擬的情況下做到這一點？

>>> from mock import MagicMock
>>> mock = MagicMock(side_effect = [1,2])
>>> mock(), mock()
(1, 2)
>>> mock()

Traceback (most recent call last):
  File "<pyshell#114>", line 1, in <module>
    mock()
  File "C:\Python27\Lib\site-packages\mock.py", line 955, in __call__
    return _mock_self._mock_call(*args, **kwargs)
  File "C:\Python27\Lib\site-packages\mock.py", line 1013, in _mock_call
    result = next(effect)
StopIteration
>>> mock.reset_mock()
>>> mock()

Traceback (most recent call last):
  ...
StopIteration
>>> mock = MagicMock(side_effect = [1,2])
>>> mock(), mock()
(1, 2)
>>> 

目的是在后續測試中重新使用模擬，但我懷疑它與生成器一樣，無法重新啟動。

因此（遲到總比沒有好）在指向正確的方向后，我查看了mock.py並發現side_effect是一個迭代器對象（一旦耗盡就無法重置）：

def __set_side_effect(self, value):
    value = _try_iter(value)
    ...

def _try_iter(obj):
    ...
    try:
        return iter(obj)
    except TypeError:
        # XXXX backwards compatibility
        # but this will blow up on first call - so maybe we should fail early?
        return obj

並且def reset_mock()沒有解決副作用。

Answer 1

正如 user2357112 評論的那樣，重新分配side_effect將解決您的問題。

>>> from mock import MagicMock
>>>
>>> lst = [1, 2]
>>> mock = MagicMock(side_effect=lst)
>>> mock(), mock()
(1, 2)
>>> mock.side_effect = lst  # <-------
>>> mock(), mock()
(1, 2)

Answer 2

請注意，從 Python 3.6 開始，您可以使用以下命令重置模擬side_effect和return_value ：

mock.reset_mock(return_value=True, side_effect=True)

從文檔：

如果您想重置 return_value 或 side_effect，則將相應的參數傳遞為 True。 子模擬和返回值模擬（如果有）也被重置。