Mock infinite generator in Python
-
26-12-2019 - |
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.
e.g. 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'>
Solution
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
Licensed under: CC-BY-SA with attribution
Not affiliated with StackOverflow