diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 46d1677..c7710a2 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -4,6 +4,7 @@ Releases Unreleased ---------- +* Fixed ``duplicate_iterators=True`` for async functions spied with ``mocker.spy``. * `#547 `_: Added ``SpyType`` for annotating ``mocker.spy`` results. * Dropped support for EOL Python 3.9. * `#147 `_: Removed handling of ``RuntimeError: stop called on unstarted patcher``, which can no longer occur in the supported Python versions. diff --git a/src/pytest_mock/plugin.py b/src/pytest_mock/plugin.py index df63fb4..37956d6 100644 --- a/src/pytest_mock/plugin.py +++ b/src/pytest_mock/plugin.py @@ -206,6 +206,12 @@ async def async_wrapper(*args, **kwargs): spy_obj.spy_exception = e raise else: + if duplicate_iterators and isinstance(r, Iterator): + r, duplicated_iterator = itertools.tee(r, 2) + spy_obj.spy_return_iter = duplicated_iterator + else: + spy_obj.spy_return_iter = None + spy_obj.spy_return = r spy_obj.spy_return_list.append(r) return r diff --git a/tests/test_pytest_mock.py b/tests/test_pytest_mock.py index c497aea..aed9894 100644 --- a/tests/test_pytest_mock.py +++ b/tests/test_pytest_mock.py @@ -666,6 +666,27 @@ async def bar(self, arg): assert result == 20 +@pytest.mark.asyncio +async def test_async_spy_return_iter_duplicates_iterator_when_enabled( + mocker: MockerFixture, +) -> None: + class Foo: + async def bar(self) -> Iterator[int]: + return iter([0, 1, 2]) + + foo = Foo() + spy = mocker.spy(foo, "bar", duplicate_iterators=True) + result = await foo.bar() + + assert list(result) == [0, 1, 2] + assert spy.spy_return is not None + assert spy.spy_return_iter is not None + assert list(spy.spy_return_iter) == [0, 1, 2] + + [return_value] = spy.spy_return_list + assert isinstance(return_value, Iterator) + + @contextmanager def assert_traceback() -> Generator[None, None, None]: """