Copy context variables from non-generator fixtures

Author: bcmills

Diff for: pytest_asyncio/plugin.py

@@ -327,18 +327,7 @@ async def setup():
         setup_task = _create_task_in_context(event_loop, setup(), context)
         result = event_loop.run_until_complete(setup_task)
 
-        # Copy the context vars set by the setup task back into the ambient
-        # context for the test.
-        context_tokens = []
-        for var in context:
-            try:
-                if var.get() is context.get(var):
-                    # Not modified by the fixture, so leave it as-is.
-                    continue
-            except LookupError:
-                pass
-            token = var.set(context.get(var))
-            context_tokens.append((var, token))
+        reset_contextvars = _apply_contextvar_changes(context)
 
         def finalizer() -> None:
             """Yield again, to finalize."""
@@ -355,38 +344,15 @@ async def async_finalizer() -> None:
 
             task = _create_task_in_context(event_loop, async_finalizer(), context)
             event_loop.run_until_complete(task)
-
-            # Since the fixture is now complete, restore any context variables
-            # it had set back to their original values.
-            while context_tokens:
-                (var, token) = context_tokens.pop()
-                var.reset(token)
+            if reset_contextvars is not None:
+                reset_contextvars()
 
         request.addfinalizer(finalizer)
         return result
 
     fixturedef.func = _asyncgen_fixture_wrapper  # type: ignore[misc]
 
 
-def _create_task_in_context(loop, coro, context):
-    """
-    Return an asyncio task that runs the coro in the specified context,
-    if possible.
-
-    This allows fixture setup and teardown to be run as separate asyncio tasks,
-    while still being able to use context-manager idioms to maintain context
-    variables and make those variables visible to test functions.
-
-    This is only fully supported on Python 3.11 and newer, as it requires
-    the API added for https://github.com/python/cpython/issues/91150.
-    On earlier versions, the returned task will use the default context instead.
-    """
-    try:
-        return loop.create_task(coro, context=context)
-    except TypeError:
-        return loop.create_task(coro)
-
-
 def _wrap_async_fixture(fixturedef: FixtureDef) -> None:
     fixture = fixturedef.func
 
@@ -403,11 +369,23 @@ async def setup():
             res = await func(**_add_kwargs(func, kwargs, event_loop, request))
             return res
 
-        # Since the fixture doesn't have a cleanup phase, if it set any context
-        # variables we don't have a good way to clear them again.
-        # Instead, treat this fixture like an asyncio.Task, which has its own
-        # independent Context that doesn't affect the caller.
-        return event_loop.run_until_complete(setup())
+        context = contextvars.copy_context()
+        setup_task = _create_task_in_context(event_loop, setup(), context)
+        result = event_loop.run_until_complete(setup_task)
+
+        # Copy the context vars modified by the setup task into the current
+        # context, and (if needed) add a finalizer to reset them.
+        #
+        # Note that this is slightly different from the behavior of a non-async
+        # fixture, which would rely on the fixture author to add a finalizer
+        # to reset the variables. In this case, the author of the fixture can't
+        # write such a finalizer because they have no way to capture the Context
+        # in which the setup function was run, so we need to do it for them.
+        reset_contextvars = _apply_contextvar_changes(context)
+        if reset_contextvars is not None:
+            request.addfinalizer(reset_contextvars)
+
+        return result
 
     fixturedef.func = _async_fixture_wrapper  # type: ignore[misc]
 
@@ -432,6 +410,57 @@ def _get_event_loop_fixture_id_for_async_fixture(
     return event_loop_fixture_id
 
 
+def _create_task_in_context(loop, coro, context):
+    """
+    Return an asyncio task that runs the coro in the specified context,
+    if possible.
+
+    This allows fixture setup and teardown to be run as separate asyncio tasks,
+    while still being able to use context-manager idioms to maintain context
+    variables and make those variables visible to test functions.
+
+    This is only fully supported on Python 3.11 and newer, as it requires
+    the API added for https://github.com/python/cpython/issues/91150.
+    On earlier versions, the returned task will use the default context instead.
+    """
+    try:
+        return loop.create_task(coro, context=context)
+    except TypeError:
+        return loop.create_task(coro)
+
+
+def _apply_contextvar_changes(
+    context: contextvars.Context,
+) -> Callable[[], None] | None:
+    """
+    Copy contextvar changes from the given context to the current context.
+
+    If any contextvars were modified by the fixture, return a finalizer that
+    will restore them.
+    """
+    context_tokens = []
+    for var in context:
+        try:
+            if var.get() is context.get(var):
+                # This variable is not modified, so leave it as-is.
+                continue
+        except LookupError:
+            # This variable isn't yet set in the current context at all.
+            pass
+        token = var.set(context.get(var))
+        context_tokens.append((var, token))
+
+    if not context_tokens:
+        return None
+
+    def restore_contextvars():
+        while context_tokens:
+            (var, token) = context_tokens.pop()
+            var.reset(token)
+
+    return restore_contextvars
+
+
 class PytestAsyncioFunction(Function):
     """Base class for all test functions managed by pytest-asyncio."""
 

Diff for: tests/async_fixtures/test_async_fixtures_contextvars.py

@@ -58,9 +58,16 @@ async def var_fixture_3(var_fixture_2):
         yield
 
 
+@pytest.fixture(scope="function")
+async def var_fixture_4(var_fixture_3, request):
+    assert _context_var.get() == "value3"
+    _context_var.set("value4")
+    # Rely on fixture teardown to reset the context var.
+
+
 @pytest.mark.asyncio
 @pytest.mark.xfail(
     sys.version_info < (3, 11), reason="requires asyncio Task context support"
 )
-async def test(var_fixture_3):
-    assert _context_var.get() == "value3"
+async def test(var_fixture_4):
+    assert _context_var.get() == "value4"