在Python中模拟无限生成器

Question

我正在尝试使用该mock库模拟无限生成器函数。（或者unittest.mock如果你有Python 3.3）

这是无限生成器的最小工作示例。如果我可以成功地模拟这个，那么我将希望能够模拟我正在使用的实际函数。

import itertools
def infinite_generator():
    thing = itertools.cycle([1, 2])
    while True:
        yield next(thing)

这是我到目前为止所尝试过的：

import mock
import itertools
mock_func = mock.MagicMock()
mock_func.__iter__.return_value = itertools.cycle([1, 2])

我想像函数mock_func一样发挥作用infinite_generator。

例如，我希望能够执行以下操作：

>>> a = mock_func()
>>> next(a)
1
>>> next(a)
2
>>> next(a)
1
>>> next(a)
2

ETC。

然而，目前next(a)返回的东西像

<MagicMock name='mock().__next__()' id='3043937712'>

Answer 1

此处省略，__iter__因为您不打算迭代mock_func对象本身：

mock_func.__iter__.return_value = itertools.cycle([1, 2])

反而：

>>> mock_func = mock.Mock()
>>> mock_func.return_value = itertools.cycle([1, 2])
>>> a = mock_func()
>>> next(a)
1
>>> next(a)
2
>>> next(a)
1
>>> next(a)
2