Mock infinite generator in Python

Question

I'm trying to mock an infinite generator function using the mock library. (Or unittest.mock if you have Python 3.3)

Here is a minimum working example of an infinite generator. If I can successfully mock this, then I will hopefully be able to mock the actual function I am using.

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

This is what I have tried so far:

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

I want mock_func to function exactly as infinite_generator functions.

eg I expect to be able to do the following:

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

etc.

However, at the moment next(a) returns things like

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

Answer 1

Leave out __iter__ here because you don't intend to iterate over the mock_func object itself:

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

Instead:

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