Repository: agronholm/anyio

Issue:
support for taskgroups and cancelscopes in fixtures
Currently documented as not supported
https://github.com/anyio/pull/76#issue-323360850

But pytest-trio has recently gained support for these https://pytest-trio.readthedocs.io/en/stable/reference.html#an-important-note-about-yield-fixtures

This could be done with `@anyio.pytest_fixture`

It would also be useful to expose this machinery for other similar tasks like asgi lifespan composition

Discussion:
agronholm: I don't think I'm going to get any extra eyes on this any time soon, so I'm merging.

Patch:
diff --git a/docs/versionhistory.rst b/docs/versionhistory.rst
index e734f37..0d19b1a 100644
--- a/docs/versionhistory.rst
+++ b/docs/versionhistory.rst
@@ -5,9 +5,11 @@ This library adheres to `Semantic Versioning 2.0 <http://semver.org/>`_.
 
 **UNRELEASED**
 
-* **BACKWARDS INCOMPATIBLE** Replaced AnyIO's own ``ExceptionGroup`` class with the PEP 654
+- **BACKWARDS INCOMPATIBLE** Replaced AnyIO's own ``ExceptionGroup`` class with the PEP 654
   ``BaseExceptionGroup`` and ``ExceptionGroup``
 - Bumped minimum version of trio to v0.19
+- Changed the pytest plugin to run both the setup and teardown phases of asynchronous generator
+  fixtures within a single task to enable use cases where a context manager straddles the ``yield``
 
 **3.5.0**
 
diff --git a/src/anyio/_backends/_asyncio.py b/src/anyio/_backends/_asyncio.py
index 2d55d8d..9ee5fe1 100644
--- a/src/anyio/_backends/_asyncio.py
+++ b/src/anyio/_backends/_asyncio.py
@@ -12,6 +12,7 @@ from asyncio import run as native_run
 from asyncio import sleep
 from asyncio.base_events import _run_until_complete_cb  # type: ignore[attr-defined]
 from collections import OrderedDict, deque
+from collections.abc import Iterable
 from concurrent.futures import Future
 from contextvars import Context, copy_context
 from dataclasses import dataclass
@@ -25,8 +26,8 @@ from socket import AddressFamily, SocketKind
 from threading import Thread
 from types import TracebackType
 from typing import (
-    IO, Any, Awaitable, Callable, Collection, ContextManager, Coroutine, Deque, Generator,
-    Iterator, Mapping, Optional, Sequence, Tuple, TypeVar, cast)
+    IO, Any, AsyncGenerator, Awaitable, Callable, Collection, ContextManager, Coroutine, Deque,
+    Generator, Iterator, Mapping, Optional, Sequence, Tuple, TypeVar, cast)
 from weakref import WeakKeyDictionary
 
 import sniffio
@@ -47,7 +48,7 @@ if sys.version_info < (3, 11):
     from exceptiongroup import BaseExceptionGroup
 
 if sys.version_info >= (3, 8):
-    def get_coro(task: asyncio.Task) -> Union[Generator, Awaitable[Any]]:
+    def get_coro(task: asyncio.Task) -> Generator | Awaitable[Any]:
         return task.get_coro()
 else:
     def get_coro(task: asyncio.Task) -> Generator | Awaitable[Any]:
@@ -1484,9 +1485,11 @@ def _create_task_info(task: asyncio.Task) -> TaskInfo:
 class TestRunner(abc.TestRunner):
     def __init__(self, debug: bool = False, use_uvloop: bool = False,
                  policy: asyncio.AbstractEventLoopPolicy | None = None):
+        self._exceptions: list[BaseException] = []
         _maybe_set_event_loop_policy(policy, use_uvloop)
         self._loop = asyncio.new_event_loop()
         self._loop.set_debug(debug)
+        self._loop.set_exception_handler(self._exception_handler)
         asyncio.set_event_loop(self._loop)
 
     def _cancel_all_tasks(self) -> None:
@@ -1505,6 +1508,19 @@ class TestRunner(abc.TestRunner):
             if task.exception() is not None:
                 raise cast(BaseException, task.exception())
 
+    def _exception_handler(self, loop: asyncio.AbstractEventLoop, context: dict[str, Any]) -> None:
+        if isinstance(context['exception'], Exception):
+            self._exceptions.append(context['exception'])
+
+    def _raise_async_exceptions(self) -> None:
+        # Re-raise any exceptions raised in asynchronous callbacks
+        if self._exceptions:
+            exceptions, self._exceptions = self._exceptions, []
+            if len(exceptions) == 1:
+                raise exceptions[0]
+            elif exceptions:
+                raise ExceptionGroup(exceptions)
+
     def close(self) -> None:
         try:
             self._cancel_all_tasks()
@@ -1513,27 +1529,51 @@ class TestRunner(abc.TestRunner):
             asyncio.set_event_loop(None)
             self._loop.close()
 
-    def call(self, func: Callable[..., Awaitable[T_Retval]],
-             *args: object, **kwargs: object) -> T_Retval:
-        def exception_handler(loop: asyncio.AbstractEventLoop, context: dict[str, Any]) -> None:
-            exceptions.append(context['exception'])
+    def run_asyncgen_fixture(self, fixture_func: Callable[..., AsyncGenerator[T_Retval, Any]],
+                             kwargs: dict[str, Any]) -> Iterable[T_Retval]:
+        async def fixture_runner() -> None:
+            agen = fixture_func(**kwargs)
+            try:
+                retval = await agen.asend(None)
+                self._raise_async_exceptions()
+            except BaseException as exc:
+                f.set_exception(exc)
+                return
+            else:
+                f.set_result(retval)
 
-        exceptions: list[BaseException] = []
-        self._loop.set_exception_handler(exception_handler)
+            await event.wait()
+            try:
+                await agen.asend(None)
+            except StopAsyncIteration:
+                pass
+            else:
+                await agen.aclose()
+                raise RuntimeError('Async generator fixture did not stop')
+
+        f = self._loop.create_future()
+        event = asyncio.Event()
+        fixture_task = self._loop.create_task(fixture_runner())
+        self._loop.run_until_complete(f)
+        yield f.result()
+        event.set()
+        self._loop.run_until_complete(fixture_task)
+        self._raise_async_exceptions()
+
+    def run_fixture(self, fixture_func: Callable[..., Coroutine[Any, Any, T_Retval]],
+                    kwargs: dict[str, Any]) -> T_Retval:
+        retval = self._loop.run_until_complete(fixture_func(**kwargs))
+        self._raise_async_exceptions()
+        return retval
+
+    def run_test(self, test_func: Callable[..., Coroutine[Any, Any, Any]],
+                 kwargs: dict[str, Any]) -> None:
         try:
-            retval: T_Retval = self._loop.run_until_complete(func(*args, **kwargs))
+            self._loop.run_until_complete(test_func(**kwargs))
         except Exception as exc:
-            retval = None  # type: ignore[assignment]
-            exceptions.append(exc)
-        finally:
-            self._loop.set_exception_handler(None)
+            self._exceptions.append(exc)
 
-        if len(exceptions) == 1:
-            raise exceptions[0]
-        elif exceptions:
-            raise ExceptionGroup(exceptions)
-
-        return retval
+        self._raise_async_exceptions()
 
 
 class AsyncIOBackend(AsyncBackend):
diff --git a/src/anyio/_backends/_trio.py b/src/anyio/_backends/_trio.py
index 0db295c..0b5c067 100644
--- a/src/anyio/_backends/_trio.py
+++ b/src/anyio/_backends/_trio.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import array
 import math
 import socket
+from collections.abc import Iterable
 from concurrent.futures import Future
 from dataclasses import dataclass
 from functools import partial
@@ -12,8 +13,8 @@ from signal import Signals
 from socket import AddressFamily, SocketKind
 from types import TracebackType
 from typing import (
-    IO, Any, Awaitable, Callable, Collection, ContextManager, Coroutine, Deque, Generic, Mapping,
-    NoReturn, Sequence, TypeVar, cast)
+    IO, TYPE_CHECKING, Any, AsyncGenerator, Awaitable, Callable, Collection, ContextManager,
+    Coroutine, Deque, Generic, Mapping, NoReturn, Sequence, TypeVar, cast)
 
 import trio.from_thread
 from outcome import Error, Outcome, Value
@@ -46,6 +47,9 @@ try:
 except ImportError:
     from trio import open_process as trio_open_process
 
+if TYPE_CHECKING:
+    from trio_typing import TaskStatus
+
 T = TypeVar('T')
 T_Retval = TypeVar('T_Retval')
 T_SockAddr = TypeVar('T_SockAddr', str, IPSockAddrType)
@@ -609,14 +613,7 @@ class TestRunner(abc.TestRunner):
     def _main_task_finished(self, outcome: object) -> None:
         self._nursery = None
 
-    def close(self) -> None:
-        if self._stop_event:
-            self._stop_event.set()
-            while self._nursery is not None:
-                self._call_queue.get()()
-
-    def call(self, func: Callable[..., Awaitable[T_Retval]],
-             *args: object, **kwargs: object) -> T_Retval:
+    def _get_nursery(self) -> trio.Nursery:
         if self._nursery is None:
             trio.lowlevel.start_guest_run(
                 self._trio_main, run_sync_soon_threadsafe=self._call_queue.put,
@@ -624,13 +621,51 @@ class TestRunner(abc.TestRunner):
             while self._nursery is None:
                 self._call_queue.get()()
 
-        self._nursery.start_soon(self._call_func, func, args, kwargs)
+        return self._nursery
+
+    def _call(self, func: Callable[..., Awaitable[T_Retval]],
+              *args: object, **kwargs: object) -> T_Retval:
+        self._get_nursery().start_soon(self._call_func, func, args, kwargs)
         while not self._result_queue:
             self._call_queue.get()()
 
         outcome = self._result_queue.pop()
         return outcome.unwrap()
 
+    def close(self) -> None:
+        if self._stop_event:
+            self._stop_event.set()
+            while self._nursery is not None:
+                self._call_queue.get()()
+
+    def run_asyncgen_fixture(self, fixture_func: Callable[..., AsyncGenerator[T_Retval, Any]],
+                             kwargs: dict[str, Any]) -> Iterable[T_Retval]:
+        async def fixture_runner(*, task_status: TaskStatus) -> None:
+            agen = fixture_func(**kwargs)
+            retval = await agen.asend(None)
+            task_status.started(retval)
+            await teardown_event.wait()
+            try:
+                await agen.asend(None)
+            except StopAsyncIteration:
+                pass
+            else:
+                await agen.aclose()
+                raise RuntimeError('Async generator fixture did not stop')
+
+        teardown_event = trio.Event()
+        fixture_value = self._call(lambda: self._get_nursery().start(fixture_runner))
+        yield fixture_value
+        teardown_event.set()
+
+    def run_fixture(self, fixture_func: Callable[..., Coroutine[Any, Any, T_Retval]],
+                    kwargs: dict[str, Any]) -> T_Retval:
+        return self._call(fixture_func, **kwargs)
+
+    def run_test(self, test_func: Callable[..., Coroutine[Any, Any, Any]],
+                 kwargs: dict[str, Any]) -> None:
+        self._call(test_func, **kwargs)
+
 
 class TrioBackend(AsyncBackend):
     @classmethod


Test patch:
diff --git a/docs/testing.rst b/docs/testing.rst
index 6e46699..6a03d95 100644
--- a/docs/testing.rst
+++ b/docs/testing.rst
@@ -150,3 +150,7 @@ shut down when that same fixture is being torn down or the test has finished run
 if no async fixtures are used, a separate test runner is created for each test. Conversely, if
 even one async fixture (scoped higher than ``function``) is shared across all tests, only one test
 runner will be created during the test session.
+
+For async generator based fixtures, the test runner spawns a task that handles both the setup and
+teardown phases to enable context-sensitive code to work properly. A common example of this is
+providing a task group as a fixture.
diff --git a/src/anyio/abc/_testing.py b/src/anyio/abc/_testing.py
index 75ed79b..320f8d4 100644
--- a/src/anyio/abc/_testing.py
+++ b/src/anyio/abc/_testing.py
@@ -2,7 +2,8 @@ from __future__ import annotations
 
 import types
 from abc import ABCMeta, abstractmethod
-from typing import Any, Awaitable, Callable, TypeVar
+from collections.abc import AsyncGenerator, Iterable
+from typing import Any, Callable, Coroutine, TypeVar
 
 _T = TypeVar("_T")
 
@@ -27,13 +28,39 @@ class TestRunner(metaclass=ABCMeta):
         """Close the event loop."""
 
     @abstractmethod
-    def call(self, func: Callable[..., Awaitable[_T]],
-             *args: object, **kwargs: dict[str, Any]) -> _T:
+    def run_asyncgen_fixture(
+        self,
+        fixture_func: Callable[..., AsyncGenerator[_T, Any]],
+        kwargs: dict[str, Any]
+    ) -> Iterable[_T]:
         """
-        Call the given function within the backend's event loop.
+        Run an async generator fixture.
 
-        :param func: a callable returning an awaitable
-        :param args: positional arguments to call ``func`` with
-        :param kwargs: keyword arguments to call ``func`` with
-        :return: the return value of ``func``
+        :param fixture_func: the fixture function
+        :param kwargs: keyword arguments to call the fixture function with
+        :return: an iterator yielding the value yielded from the async generator
+        """
+
+    @abstractmethod
+    def run_fixture(
+        self,
+        fixture_func: Callable[..., Coroutine[Any, Any, _T]],
+        kwargs: dict[str, Any]
+    ) -> _T:
+        """
+        Run an async fixture.
+
+        :param fixture_func: the fixture function
+        :param kwargs: keyword arguments to call the fixture function with
+        :return: the return value of the fixture function
+        """
+
+    @abstractmethod
+    def run_test(self, test_func: Callable[..., Coroutine[Any, Any, Any]],
+                 kwargs: dict[str, Any]) -> None:
+        """
+        Run an async test function.
+
+        :param test_func: the test function
+        :param kwargs: keyword arguments to call the test function with
         """
diff --git a/src/anyio/pytest_plugin.py b/src/anyio/pytest_plugin.py
index ac4af75..87673ca 100644
--- a/src/anyio/pytest_plugin.py
+++ b/src/anyio/pytest_plugin.py
@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, Dict, Iterator, Tuple, cast
 
 import pytest
 import sniffio
+from _pytest.fixtures import FixtureRequest
 
 from ._core._eventloop import get_all_backends, get_async_backend
 from .abc import TestRunner
@@ -55,7 +56,7 @@ def pytest_configure(config: Config) -> None:
                                        'asynchronously via anyio.')
 
 
-def pytest_fixture_setup(fixturedef: Any, request: Any) -> None:
+def pytest_fixture_setup(fixturedef: Any, request: FixtureRequest) -> None:
     def wrapper(*args, anyio_backend, **kwargs):  # type: ignore[no-untyped-def]
         backend_name, backend_options = extract_backend_and_options(anyio_backend)
         if has_backend_arg:
@@ -63,23 +64,9 @@ def pytest_fixture_setup(fixturedef: Any, request: Any) -> None:
 
         with get_runner(backend_name, backend_options) as runner:
             if isasyncgenfunction(func):
-                gen = func(*args, **kwargs)
-                try:
-                    value = runner.call(gen.asend, None)
-                except StopAsyncIteration:
-                    raise RuntimeError('Async generator did not yield')
-
-                yield value
-
-                try:
-                    runner.call(gen.asend, None)
-                except StopAsyncIteration:
-                    pass
-                else:
-                    runner.call(gen.aclose)
-                    raise RuntimeError('Async generator fixture did not stop')
+                yield from runner.run_asyncgen_fixture(func, kwargs)
             else:
-                yield runner.call(func, *args, **kwargs)
+                yield runner.run_fixture(func, kwargs)
 
     # Only apply this to coroutine functions and async generator functions in requests that involve
     # the anyio_backend fixture
@@ -107,7 +94,7 @@ def pytest_pycollect_makeitem(collector: Any, name: Any, obj: Any) -> None:
 def pytest_pyfunc_call(pyfuncitem: Any) -> bool | None:
     def run_with_hypothesis(**kwargs: Any) -> None:
         with get_runner(backend_name, backend_options) as runner:
-            runner.call(original_func, **kwargs)
+            runner.run_test(original_func, kwargs)
 
     backend = pyfuncitem.funcargs.get('anyio_backend')
     if backend:
@@ -126,14 +113,14 @@ def pytest_pyfunc_call(pyfuncitem: Any) -> bool | None:
             funcargs = pyfuncitem.funcargs
             testargs = {arg: funcargs[arg] for arg in pyfuncitem._fixtureinfo.argnames}
             with get_runner(backend_name, backend_options) as runner:
-                runner.call(pyfuncitem.obj, **testargs)
+                runner.run_test(pyfuncitem.obj, testargs)
 
             return True
 
     return None
 
 
-@pytest.fixture(params=get_all_backends())
+@pytest.fixture(scope="module", params=get_all_backends())
 def anyio_backend(request: Any) -> Any:
     return request.param
 
diff --git a/tests/test_pytest_plugin.py b/tests/test_pytest_plugin.py
index f3b0e0c..56049a3 100644
--- a/tests/test_pytest_plugin.py
+++ b/tests/test_pytest_plugin.py
@@ -33,8 +33,6 @@ def test_plugin(testdir: Pytester) -> None:
 
     testdir.makepyfile(
         """
-        import asyncio
-
         import pytest
         import sniffio
         from hypothesis import strategies, given
@@ -83,7 +81,7 @@ def test_asyncio(testdir: Pytester) -> None:
             def callback():
                 raise RuntimeError('failing fixture setup')
 
-            asyncio.get_event_loop().call_soon(callback)
+            asyncio.get_running_loop().call_soon(callback)
             await asyncio.sleep(0)
             yield None
 
@@ -93,7 +91,7 @@ def test_asyncio(testdir: Pytester) -> None:
                 raise RuntimeError('failing fixture teardown')
 
             yield None
-            asyncio.get_event_loop().call_soon(callback)
+            asyncio.get_running_loop().call_soon(callback)
             await asyncio.sleep(0)
         """
     )
@@ -124,7 +122,7 @@ def test_asyncio(testdir: Pytester) -> None:
                 raise Exception('foo')
 
             started = False
-            asyncio.get_event_loop().call_soon(callback)
+            asyncio.get_running_loop().call_soon(callback)
             await asyncio.sleep(0)
             assert started
 
@@ -204,6 +202,53 @@ def test_cancel_scope_in_asyncgen_fixture(testdir: Pytester) -> None:
     result.assert_outcomes(passed=len(get_all_backends()))
 
 
+def test_module_scoped_task_group_fixture(testdir: Pytester) -> None:
+    testdir.makeconftest(
+        """
+        import pytest
+
+        from anyio import CancelScope, create_memory_object_stream, create_task_group
+
+
+        @pytest.fixture(scope="module")
+        async def task_group():
+            async with create_task_group() as tg:
+                yield tg
+
+
+        @pytest.fixture
+        async def streams(task_group):
+            async def echo_messages(*, task_status):
+                with CancelScope() as cancel_scope:
+                    task_status.started(cancel_scope)
+                    async for obj in receive1:
+                        await send2.send(obj)
+
+            send1, receive1 = create_memory_object_stream()
+            send2, receive2 = create_memory_object_stream()
+            cancel_scope = await task_group.start(echo_messages)
+            yield send1, receive2
+            cancel_scope.cancel()
+        """
+    )
+
+    testdir.makepyfile(
+        """
+        import pytest
+
+
+        @pytest.mark.anyio
+        async def test_task_group(streams):
+            send1, receive2 = streams
+            await send1.send("hello")
+            assert await receive2.receive() == "hello"
+        """
+    )
+
+    result = testdir.runpytest_subprocess(*pytest_args)
+    result.assert_outcomes(passed=len(get_all_backends()))
+
+
 def test_hypothesis_module_mark(testdir: Pytester) -> None:
     testdir.makepyfile(
         """