CODEWITHSUNDEEP
pythonunit-testingmocking
TinyTheBrontosaurus
TinyTheBrontosaurus

Reputation: 4810

How to reset mock iterator inside a for loop?

Similar to this question, but the answers in there are not adequate for what I'm doing.

I'm trying to test a method like this:

import mock


def stack_overflow_desired_output():
    print_a_few_times(['upvote', 'this', 'question!'])


def stack_overflow_mocked():
    the_mock = mock.Mock()
    the_mock.__iter__ = mock.Mock(return_value=iter(["upvote", "this", "question"]))
    print_a_few_times(the_mock)


def print_a_few_times(fancy_object):
    for x in [1, 2, 3]:
        for y in fancy_object:
            print("{}.{}".format(x, y))

When I call stack_overflow_desired_output() I get this:

1.upvote
1.this
1.question!
2.upvote
2.this
2.question!
3.upvote
3.this
3.question!

But when I call stack_overflow_mocked(), I get only get this:

1.upvote
1.this
1.question!

Is there a way to have that iterator reset itself whenever it's exhausted at the end of a for loop? Putting the reset inside of the print_a_few_times, function, as described in the aforementioned question, will be invasive.

Upvotes: 2

Views: 857

Answers (1)

chepner
chepner

Reputation: 531315

Wrap your mock object around the __iter__ method of an actual list.

def stack_overflow_mocked():
    the_mock = mock.Mock()
    the_mock.__iter__ = mock.Mock(wraps=["upvote", "this", "question"].__iter__)
    print_a_few_times(the_mock)

Upvotes: 1

Related Questions