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

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

Answer 3

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