Mocking Generator Methods in Python

Another mock recipe, this one for mocking generator methods.

A Python generator is a function or method that uses the yield statement to return a series of values when iterated over (there are also generator expressions and more advanced uses of generators, but we aren't concerned about them here. A very good introduction to generators and how powerful they are is: Generator Tricks for Systems Programmers.).

A generator method / function is called to return the generator object. It is the generator object that is then iterated over. The protocol method for iteration is __iter__, so we can mock this using a MagicMock.

Here's an example class with an "iter" method implemented as a generator:

>>> class Foo(object):
...     def iter(self):
...         for i in [1, 2, 3]:
...             yield i
...
>>> foo = Foo()
>>> list(foo.iter())
[1, 2, 3]

How would we mock this class, and in particular its "iter" method?

To configure the values returned from the iteration (implicit in the call to list), we need to configure the iterator returned by the call to foo.iter().

>>> values = [1, 2, 3]
>>> mock_foo = MagicMock()
>>> iterable = mock_foo.iter.return_value
>>> iterator = iter(values)
>>> iterable.__iter__.return_value = iterator
>>> list(mock_foo.iter())
[1, 2, 3]

The above example is done step-by-step. The shorter version is:

>>> mock_foo = MagicMock()
>>> mock_foo.iter.return_value.__iter__.return_value = iter([1, 2, 3])
>>> list(mock_foo.iter())
[1, 2, 3]

This is now also in the docs on the mock examples page. There's now quite a collection of useful mock recipes there, so even if you're an experienced mock user it is worth a browse.

Source: http://www.voidspace.org.uk/python/weblog/arch_d7_2011_06_11.shtml