From 97c682f48c60b60c07e977e6cd74e844d43ccd68 Mon Sep 17 00:00:00 2001 From: "Bryan C. Mills" Date: Mon, 9 Dec 2024 09:37:46 -0500 Subject: [PATCH] Copy context variables from non-generator fixtures --- pytest_asyncio/plugin.py | 113 +++++++++++------- .../test_async_fixtures_contextvars.py | 11 +- 2 files changed, 80 insertions(+), 44 deletions(-) diff --git a/pytest_asyncio/plugin.py b/pytest_asyncio/plugin.py index 062437b7..e1ce5c9d 100644 --- a/pytest_asyncio/plugin.py +++ b/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,12 +344,8 @@ 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 @@ -368,25 +353,6 @@ async def async_finalizer() -> None: 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 --git a/tests/async_fixtures/test_async_fixtures_contextvars.py b/tests/async_fixtures/test_async_fixtures_contextvars.py index bddab379..8d67db49 100644 --- a/tests/async_fixtures/test_async_fixtures_contextvars.py +++ b/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"