Python Mock Object Iterator Not Able To Iterate Multiple Times

Question

I am using the mock iterator functionality to return an iterator for unit tests. In the code that I am testing, the I loop through the object multiple times, but it doesn't seem to be working and only works the first time.

self.mock_scene.bpyscene.objects.__iter__ = mock.Mock(return_value=iter([mock_lamp_object, mock_lamp_object]))

Answer 1

You can use the side_effect parameter of Mock to override the __next__ attribute of the class you're testing instead.

According to the documentation :

If side_effect is an iterable then each call to the mock will return the next value from the iterable.

Example:

from unittest.mock import Mock

class Iterable:
    def __iter__(self):
        return self

Iterable.__next__ = Mock(side_effect=[1, 2, 3])

for i in Iterable():
    print(i)

This outputs:

1
2
3

Answer 2

Figured it out. You can only iterate once through a mock iterator, after that it is exhausted. To get around this use MagicMock and its iterator, which can be used as many times as you like.