Skip to content

Conversation

vokracko
Copy link

@vokracko vokracko commented Sep 1, 2025

Recently I've run into situation where I wanted to spy on a method returning Iterator.
That wasn't possible because the iterator is already consumed by the time you try to assert it.

In this PR I add a new attribute, spy_return_iter storing a duplicate of the returned iterator (leveraging itertools.tee).

Copy link
Member

@nicoddemus nicoddemus left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the contribution @vokracko!

Left some minor comments, please take a look.

assert spy.spy_return is not None
assert spy.spy_return_iter is not None
assert list(spy.spy_return_iter) == result
assert spy.spy_return_list == [ANY]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this work?

Suggested change
assert spy.spy_return_list == [ANY]
assert isinstance(spy.spy_return_list[0], Iterator)

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added in e6c26b2

Comment on lines +583 to +602
def test_spy_return_iter_unset_in_last_call(mocker: MockerFixture) -> None:
class Foo:
iterables = [
(i for i in range(3)),
[3, 4, 5],
]

def bar(self) -> Iterable[int]:
return self.iterables.pop(0)

foo = Foo()
spy = mocker.spy(foo, "bar")
result_iterator = list(foo.bar())

assert result_iterator == [0, 1, 2]
assert list(spy.spy_return_iter) == result_iterator

result_iterable = foo.bar()
assert result_iterable == [3, 4, 5]
assert spy.spy_return_iter is None
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor suggestion to make this clearer:

Suggested change
def test_spy_return_iter_unset_in_last_call(mocker: MockerFixture) -> None:
class Foo:
iterables = [
(i for i in range(3)),
[3, 4, 5],
]
def bar(self) -> Iterable[int]:
return self.iterables.pop(0)
foo = Foo()
spy = mocker.spy(foo, "bar")
result_iterator = list(foo.bar())
assert result_iterator == [0, 1, 2]
assert list(spy.spy_return_iter) == result_iterator
result_iterable = foo.bar()
assert result_iterable == [3, 4, 5]
assert spy.spy_return_iter is None
def test_spy_return_iter_resets(mocker: MockerFixture) -> None:
class Foo:
iterables = [
(i for i in range(3)),
99,
]
def bar(self) -> Iterable[int]:
return self.iterables.pop(0)
foo = Foo()
spy = mocker.spy(foo, "bar")
result_iterator = list(foo.bar())
assert result_iterator == [0, 1, 2]
assert list(spy.spy_return_iter) == result_iterator
assert foo.bar() == 99
assert spy.spy_return_iter is None

Mixing iterator and iterables on this test is not necessary for the logic and adds a small bit of confusion.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants