How to mock generators with mock.patch
Wims answer:
mock_f.return_value = iter(items)
works as long as your mock gets called only once. In unit testing, we may often want to call a function or method multiple times with different arguments. That will fail in this case, because on the first call the iterator will be exhausted such that on the second call it will immediately raise a StopIteration
exception. With Alexandre Paes' answer I was getting an AttributeError: 'function' object has no attribute '__iter__'
when my mock was coming from unittest.mock.patch
.
As an alternative, we can create a “fake” iterator and assign that as a side_effect
:
@unittest.mock.patch("mymod.my_generator", autospec=True):
def test_my_func(mm):
from mymod import my_func
def fake():
yield from [items]
mm.side_effect = fake
my_func() # which calls mymod.my_generator
my_func() # subsequent calls work without unwanted memory from first call
I have another approach:
mock_f.__iter__.return_value = [items]
This way you really mock the iterator returned value.
This approach works even when you are mocking complex objects which are iterables and have methods (my case).
I tried the chosen answer but didtn't work in my case, only worked when I mocked the way I explained
Change this line:
mock_f.iter.return_value = items
To this:
mock_f.return_value = iter(items)