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

等等

但是，目前 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