Mocking iterators

A colleague wanted to mock a Journal object which both has callable methods and works as an iterator itself. So it works like this:

j = Journal()
j.log_level(Journal.INFO)
for line in j:
   print(line)

We mocked it like this, to be able to pass an actual list of expected values the function will iterate over:

import mock

mock_journal = mock.Mock()
mock_journal.__next__ = mock.Mock(side_effect=[1,2,3,4])
mock_journal.__iter__ = mock.Mock(return_value=mock_journal)


for i in mock_journal:
    print(i)
# I don't call any methods in mock_journal, but I could,
# and could then assert they were called.

So mock_journal is both a mock proper, where methods can be called (and then asserted on), and an iterable, which when called repeatedly will yield elements of the __next__ side_effect.

Leave a Reply

This site uses Akismet to reduce spam. Learn how your comment data is processed.