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

score 44 · Accepted Answer · answered Mar 22 '16 at 22:00

44

Change this line:

mock_f.iter.return_value = items

To this:

mock_f.return_value = iter(items)

answered Mar 22 '16 at 22:00

wim

338,267
99
616
750

3

Be careful with this! When you call the method multiple times it will be traversing through the items in the iterable. – D Hudson Nov 20 '19 at 14:32

Alexandre Paes · Answer 2 · 2019-12-19T11:15:02.240

5

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

edited Dec 19 '19 at 11:15

answered Dec 18 '19 at 13:57

Alexandre Paes

121
2
4

2

Nice approach, but it is good pratice to explain why your answer is different and/or better than the one already excepted. – nicolasassi Dec 18 '19 at 14:00

score 5 · Answer 3 · answered Nov 26 '20 at 09:21

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

What does `autosepc` and `side_effec` does here? It feels like magic. It works. — buhtz, Jan 31 '22 at 22:34
@buhtz `autospec` makes sure that the mock requires the same spec (input args) as the original. The `side_effect` is a function or generator called every time the mock is called. See the unittest.mock docs for details. — gerrit, Feb 01 '22 at 12:41

How to mock generators with mock.patch

3 Answers3

Linked