ContextVar not propagated from fixture to test

[beedub]

I'm trying to figure out why this code sample isn't working. I'm not experienced with asyncio yet, so this is very likely a simple misunderstanding on my part.

from _contextvars import ContextVar

import pytest

blah = ContextVar("blah")


@pytest.fixture
async def my_context_var():
    blah.set("hello")
    assert blah.get() == "hello"
    yield blah


@pytest.mark.asyncio
async def test_blah(my_context_var):
    assert my_context_var.get() == "hello" # this fails

[kalefranz]

I've also hit this. I think what's happening is explained by

pytest-asyncio/pytest_asyncio/plugin.py

Line 97 in 2c8023d

return loop.run_until_complete(setup())

along with understanding of https://www.python.org/dev/peps/pep-0550. The setup() function there executes the code within the coroutine fixture. Using loop.run_until_complete(setup()) means that context changes won't "leak" outside of the execution context within the coroutine.

One option is to try to avoid using a coroutine fixture when you need to set context. So for example

import asyncio
from contextvars import ContextVar
import pytest

blah = ContextVar("blah")


async def set_context():
    blah.set("hello")
    assert blah.get() == "hello"


@pytest.fixture
def my_context_var():
    asyncio.run(set_context())
    assert blah.get() == "hello"
    yield blah


@pytest.mark.asyncio
async def test_blah(my_context_var):
    assert my_context_var.get() == "hello"

[Andrei-Pozolotin]

+1

[nolar]

A note on the workaround above: the set_context() coroutine will run in its own context, and its contextvar will not get into the surrounding/sibling coroutines. The example above fails with:

__________________________________________________________________________ ERROR at setup of test_blah ___________________________________________________________________________

    @pytest.fixture
    def my_context_var():
        asyncio.run(set_context())
>       assert blah.get() == "hello"
E       LookupError: <ContextVar name='blah' at 0x102ebb900>

A much easier solution is to set the contextvar directly in the sync fixture — note, not async fixture, but sync:

from contextvars import ContextVar
import pytest

blah = ContextVar("blah")


@pytest.fixture
def my_context_var():  # << SYNC!!!
    blah.set("hello")
    assert blah.get() == "hello"
    yield blah


@pytest.mark.asyncio
async def test_blah(my_context_var):
    assert my_context_var.get() == "hello"

I use this approach in Kopf actively, and it works:

• https://github.com/zalando-incubator/kopf/blob/0.26/tests/posting/conftest.py
• https://github.com/zalando-incubator/kopf/blob/0.26/tests/conftest.py#L322-L344
• https://github.com/zalando-incubator/kopf/blob/0.26/tests/authentication/test_credentials.py#L91-L95
• And basically everything you can find there with _var.set(

[dimaqq]

The cause appears to be that the async fixtures are ran in own Tasks.

To augment @nolar 's solution, I'd recommend this:

@pytest.fixture
def fake_db(database):
    token = DB.set(database)
    yield
    DB.reset(token)

@pytest.fixture
async def database(...):
    ...

[butla]

I have async functions that set context vars with objects that should be used in the entire app (I use a file with context vars as an IoC container). I wanted to run that code in a fixture, so I get nice setup and teardown, but sadly, I can't run that in a sync fixture. So the workarounds don't work for me.

Just letting anybody interested know :)

[nikicat]

This issue could be fixed easily in Python 3.11, but even for 3.10 this could be fixed like this

[rinarakaki]

@nikicat Can you show us how to fix this issue in Python 3.11 please?

[dimaqq]

@rnarkk I believe that #127 (comment) still works.

The essence of Nikolay's suggestion is to make a custom event_loop fixture which sets a custom task factory on the event loop. The fixture name shadows the default fixture from pytest-asyncio and thus the custom fixture gets automatically used.

[seifertm]

#620 introduced the asyncio_event_loop mark, which provides a class-scoped or module-scoped event loop when a class or module is marked, respectively. #646 plans to extend the mark so it can be applied to test coroutines as well.

Under the hood, the mark creates an event loop and attaches it to the marked pytest Item. Tests and fixtures run in the same event loop when the asyncio_event_loop mark is used. Unfortunately, tests and fixtures are run using loop.run_until_complete which still prevents context variables from being propagated from fixtures to tests.

The next step is to make asyncio_event_loop attach an asyncio.Runner, instead of an event loop to the marked pytest Item. This allows replacing loop.run_until_complete with asyncio.Runner.run, which preserves contextvars. We will also need a backport for Python 3.10 and earlier.

There's still work to be done, but I wanted to give a progress update and let you know that this issue hasn't been forgotten.

[mjvankampen]

@seifertm any update on this. #646 was replaced by #657 and #667. I'm not sure what the best practice is to circumvent this issue with the current version of pytest-asyncio implementation.

[pkucmus]

This is very troublesome with encode/databases >=0.8.0 (related). They keep state in context vars which results in a lost transaction reference - the DB connection and a unit test rollback transaction is made in a fixture, then in the test the transaction is gone breaking an assert in encode/databases.

I was able to get around it with the help of the comment from above and ended up with this - this solves my issue with encode/databases:

def task_factory(loop, coro, context=None):
    stack = traceback.extract_stack()
    for frame in stack[-2::-1]:
        package_name = Path(frame.filename).parts[-2]
        if package_name != "asyncio":
            if package_name == "pytest_asyncio":
                # This function was called from pytest_asyncio, use shared context
                break
            else:
                # This function was called from somewhere else, create context copy
                context = None
            break
    return asyncio.Task(coro, loop=loop, context=context)


@pytest.fixture(scope="session")
def event_loop():
    loop = asyncio.get_event_loop_policy().new_event_loop()
    context = contextvars.copy_context()
    loop.set_task_factory(functools.partial(task_factory, context=context))
    asyncio.set_event_loop(loop)
    return loop

It complains that "Replacing the event_loop fixture with a custom implementation is deprecated and will lead to errors in the future." but nothing else I tried seems to work, the context vars are always gone. I thought this would work best:

class CustomEventLoop(asyncio.SelectorEventLoop):
    _task_factory = functools.partial(task_factory, context=contextvars.copy_context())

    # also tried to overload create_task, still no luck 
    def create_task(self, coro, *, name=None, context=None):
        context = contextvars.copy_context()
        print("Creating task with context", context)
        return super().create_task(coro, name=name, context=context)


class CustomEventLoopPolicy(asyncio.DefaultEventLoopPolicy):
    _loop_factory = CustomEventLoop


@pytest.fixture(scope="session")
def event_loop_policy():
    policy = CustomEventLoopPolicy()
    yield policy
    policy.get_event_loop().close()

but neither with _task_factory nor create_task could I force pytest to retain the context vars. The deprecated event_loop fixture is the only one that seems to work.

[pkucmus]

OK, I was able to make it stop complaining about the deprecation with the following:

import asyncio
import functools

import traceback
from pathlib import Path

def task_factory(loop, coro, context=None):
    stack = traceback.extract_stack()
    for frame in stack[-2::-1]:
        package_name = Path(frame.filename).parts[-2]
        if package_name != "asyncio":
            if package_name == "pytest_asyncio":
                # This function was called from pytest_asyncio, use shared context
                break
            else:
                # This function was called from somewhere else, create context copy
                context = None
            break
    return asyncio.Task(coro, loop=loop, context=context)


class CustomEventLoopPolicy(asyncio.DefaultEventLoopPolicy):

    def __init__(self, context) -> None:
        super().__init__()
        self.context = context

    def new_event_loop(self):
        loop = self._loop_factory()
        loop.set_task_factory(
            functools.partial(task_factory, context=self.context)
        )
        return loop


@pytest.fixture(scope="session")
def event_loop_policy():
    policy = CustomEventLoopPolicy(contextvars.copy_context())
    yield policy
    policy.get_event_loop().close()