在Python中模拟无限生成器

Question

我正在尝试使用 mock 库来模拟一个无限生成器函数。（如果你用的是 Python 3.3，可以用 unittest.mock）

这里有一个最简单的无限生成器的例子。如果我能成功模拟这个生成器，那么我希望能模拟我正在使用的实际函数。

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