Mocking __iter__ with a magic mock in Python

I just ran into a tiny little problem and thought I’d share the solution.

I was writing some tests in which I was mocking a dependency using Mock. My test code was roughly equivalent to the following (based on an example in the Mock documentation on magic methods):

>>> sub1 = Mock()
>>> sub1.value = 1
>>> sub2 = Mock()
>>> sub2.value = 2
>>> m = Mock()
>>> m.__iter__ = Mock(return_value = iter([sub1, sub2]))

The point here was to mock the __iter__ method to be able to use the main mock in a list comprehension statement. The odd thing was that my tests failed, and I quickly found that of two list comprehension statements, the second would return an empty result:

>>> [x.value for x in m if x.value == 1]
[1]

>>> [x.value for x in m if x.value == 2]
[]

An educated guess is that the first iteration consumes all elements and leaves the iterator exhausted! As it turns out, this is clearly stated in the Python documentation about iterator objects:

One notable exception is code which attempts multiple iteration passes. […] Attempting this with an iterator will just return the same exhausted iterator object used in the previous iteration pass, making it appear like an empty container.

To solve the problem, I had to use the MagicMock class instead. The documentation says:

In 0.8 the iter also gained special handling implemented with a side effect. The return value of MagicMock.iter can be any iterable object and isn’t required to be an iterator

…and gives an example similar to the following:

m = MagicMock()
m.__iter__.return_value = [sub1, sub2]

Putting it to the test:

>>> [x.value for x in m if x.value == 1]
[1]

>>> [x.value for x in m if x.value == 2]
[2]

It works! :-) A final note: It’s really 0.8 that is required; I tried with 0.8.0beta3 (which I had installed), and it didn’t work. An upgrade was needed!