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)