How to mock generators with mock.patch

Question

I have gone through the page https://docs.python.org/3/library/unittest.mock-examples.html and i see that they have listed an example on how to mock generators

I have a code where i call a generator to give me a set of values that i save as a dictionary. I want to mock the calls to this generator in my unit test.

I have written the following code and it does not work.

Where am i going wrong?

In [7]: items = [(1,'a'),(2,'a'),(3,'a')]  In [18]: def f():     print "here"     for i in [1,2,3]:         yield i,'a'  In [8]: def call_f():    ...:     my_dict = dict(f())    ...:     print my_dict[1]    ...:   In [9]: call_f() "here" a  In [10]: import mock   In [18]: def test_call_f():     with mock.patch('__main__.f') as mock_f:         mock_f.iter.return_value = items         call_f()    ....:   In [19]: test_call_f() --------------------------------------------------------------------------- KeyError                                  Traceback (most recent call last) <ipython-input-19-33ca65a4f3eb> in <module>() ----> 1 test_call_f()  <ipython-input-18-92ff5f1363c8> in test_call_f()       2     with mock.patch('__main__.f') as mock_f:       3         mock_f.iter.return_value = items ----> 4         call_f()  <ipython-input-8-a5cff08ebf69> in call_f()       1 def call_f():       2     my_dict = dict(f()) ----> 3     print my_dict[1]  KeyError: 1 

Answer 1

Change this line:

mock_f.iter.return_value = items 

To this:

mock_f.return_value = iter(items) 

Answer 2

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