asyncio: #8 consider contextvars

Make it so everything gets executed in the same asyncio context

Author: delfick

README.rst

@@ -21,6 +21,9 @@ generator fixtures.
 Changelog
 ---------
 
+0.8.0 - TBD
+    * Provide simple support for tests being aware of asyncio.Context
+
 0.7.2 - 1 October 2023
     * Timeouts don't take affect if the debugger is active
 

alt_pytest_asyncio/async_converters.py

@@ -14,19 +14,78 @@
 from functools import wraps
 
 
-def convert_fixtures(fixturedef, request, node):
+def convert_fixtures(ctx, fixturedef, request, node):
     """Used to replace async fixtures"""
     if not hasattr(fixturedef, "func"):
         return
 
+    if hasattr(fixturedef.func, "__alt_asyncio_pytest_converted__"):
+        return
+
     if inspect.iscoroutinefunction(fixturedef.func):
-        convert_async_coroutine_fixture(fixturedef, request, node)
+        convert_async_coroutine_fixture(ctx, fixturedef, request, node)
+        fixturedef.func.__alt_asyncio_pytest_converted__ = True
 
     elif inspect.isasyncgenfunction(fixturedef.func):
-        convert_async_gen_fixture(fixturedef, request, node)
+        convert_async_gen_fixture(ctx, fixturedef, request, node)
+        fixturedef.func.__alt_asyncio_pytest_converted__ = True
+
+    elif inspect.isgeneratorfunction(fixturedef.func):
+        convert_sync_gen_fixture(ctx, fixturedef)
+        fixturedef.func.__alt_asyncio_pytest_converted__ = True
+
+    else:
+        convert_sync_fixture(ctx, fixturedef)
+        fixturedef.func.__alt_asyncio_pytest_converted__ = True
+
+
+def convert_sync_fixture(ctx, fixturedef):
+    """
+    Used to make sure a non-async fixture is run in our
+    asyncio contextvars
+    """
+    original = fixturedef.func
+
+    @wraps(original)
+    def run_fixture(*args, **kwargs):
+        try:
+            ctx.run(lambda: None)
+            run = lambda func, *a, **kw: ctx.run(func, *a, **kw)
+        except RuntimeError:
+            run = lambda func, *a, **kw: func(*a, **kw)
 
+        return run(original, *args, **kwargs)
 
-def converted_async_test(test_tasks, func, timeout, *args, **kwargs):
+    fixturedef.func = run_fixture
+
+
+def convert_sync_gen_fixture(ctx, fixturedef):
+    """
+    Used to make sure a non-async generator fixture is run in our
+    asyncio contextvars
+    """
+    original = fixturedef.func
+
+    @wraps(original)
+    def run_fixture(*args, **kwargs):
+        try:
+            ctx.run(lambda: None)
+            run = lambda func, *a, **kw: ctx.run(func, *a, **kw)
+        except RuntimeError:
+            run = lambda func, *a, **kw: func(*a, **kw)
+
+        cm = original(*args, **kwargs)
+        value = run(cm.__next__)
+        try:
+            yield value
+            run(cm.__next__)
+        except StopIteration:
+            pass
+
+    fixturedef.func = run_fixture
+
+
+def converted_async_test(ctx, test_tasks, func, timeout, *args, **kwargs):
     """Used to replace async tests"""
     __tracebackhide__ = True
 
@@ -39,7 +98,7 @@ def look_at_task(t):
     info["func"] = func
 
     return _run_and_raise(
-        loop, info, func, async_runner(func, timeout, info, args, kwargs), look_at_task
+        ctx, loop, info, func, async_runner(func, timeout, info, args, kwargs), look_at_task
     )
 
 
@@ -83,15 +142,15 @@ def raise_error():
         raise_error()
 
 
-def _run_and_raise(loop, info, func, coro, look_at_task=None):
+def _run_and_raise(ctx, loop, info, func, coro, look_at_task=None):
     __tracebackhide__ = True
 
     def silent_done_task(res):
         if res.cancelled():
             pass
         res.exception()
 
-    task = loop.create_task(coro)
+    task = loop.create_task(coro, context=ctx)
     task.add_done_callback(silent_done_task)
 
     if look_at_task:
@@ -148,7 +207,7 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-def convert_async_coroutine_fixture(fixturedef, request, node):
+def convert_async_coroutine_fixture(ctx, fixturedef, request, node):
     """
     Run our async fixture in our event loop and capture the error from
     inside the loop.
@@ -161,12 +220,14 @@ def override(extra, *args, **kwargs):
 
         info = {"func": func}
         loop = asyncio.get_event_loop_policy().get_event_loop()
-        return _run_and_raise(loop, info, func, async_runner(func, timeout, info, args, kwargs))
+        return _run_and_raise(
+            ctx, loop, info, func, async_runner(func, timeout, info, args, kwargs)
+        )
 
     fixturedef.func = _wrap(fixturedef, [], override, func)
 
 
-def convert_async_gen_fixture(fixturedef, request, node):
+def convert_async_gen_fixture(ctx, fixturedef, request, node):
     """
     Return the yield'd value from the generator and ensure the generator is
     finished.
@@ -204,11 +265,12 @@ async def async_finalizer():
                 else:
                     info["e"] = ValueError("Async generator fixture should only yield once")
 
-            _run_and_raise(loop, info, generator, async_finalizer())
+            _run_and_raise(ctx, loop, info, generator, async_finalizer())
 
         request.addfinalizer(finalizer)
 
         return _run_and_raise(
+            ctx,
             loop,
             info,
             generator,

alt_pytest_asyncio/plugin.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextvars
 import inspect
 import sys
 from collections import defaultdict
@@ -19,7 +20,9 @@ def __init__(self, loop=None):
             self.own_loop = True
             loop = asyncio.new_event_loop()
             asyncio.set_event_loop(loop)
+
         self.loop = loop
+        self.ctx = contextvars.copy_context()
 
     def pytest_configure(self, config):
         """Register our timeout marker which is used to signify async timeouts"""
@@ -57,7 +60,7 @@ def pytest_sessionfinish(self, session, exitstatus):
     @pytest.hookimpl(tryfirst=True, hookwrapper=True)
     def pytest_fixture_setup(self, fixturedef, request):
         """Convert async fixtures to sync fixtures"""
-        convert_fixtures(fixturedef, request, request.node)
+        convert_fixtures(self.ctx, fixturedef, request, request.node)
         yield
 
     @pytest.hookimpl(tryfirst=True, hookwrapper=True)
@@ -75,7 +78,25 @@ def pytest_pyfunc_call(self, pyfuncitem):
                 )
 
             o = pyfuncitem.obj
-            pyfuncitem.obj = wraps(o)(partial(converted_async_test, self.test_tasks, o, timeout))
+            pyfuncitem.obj = wraps(o)(
+                partial(converted_async_test, self.ctx, self.test_tasks, o, timeout)
+            )
+        else:
+
+            original = pyfuncitem.obj
+
+            @wraps(original)
+            def run_obj(*args, **kwargs):
+                try:
+                    self.ctx.run(lambda: None)
+                    run = lambda func, *a, **kw: self.ctx.run(func, *a, **kw)
+                except RuntimeError:
+                    run = lambda func, *a, **kw: func(*a, **kw)
+
+                run(original, *args, **kwargs)
+
+            pyfuncitem.obj = run_obj
+
         yield
 
 

pyproject.toml

@@ -25,6 +25,7 @@ tests = [
     "nest-asyncio==1.0.0",
     "noseOfYeti[black]==2.4.1",
     "pytest==7.3.0",
+    "pytest-order==1.2.1"
 ]
 
 [project.entry-points.pytest11]

tests/code_contextvars/__init__.py

@@ -0,0 +1,22 @@
+import string
+from contextvars import ContextVar
+
+allvars = {}
+
+
+class Empty:
+    pass
+
+
+def assertVarsEmpty(excluding=None):
+    for letter, var in allvars.items():
+        if not excluding or letter not in excluding:
+            assert var.get(Empty) is Empty
+
+
+for letter in string.ascii_letters:
+    var = ContextVar(letter)
+    locals()[letter] = var
+    allvars[letter] = var
+
+__all__ = ["allvars", "Empty", "assertVarsEmpty"] + sorted(allvars)

tests/code_contextvars/contextvars_for_test.py

@@ -1,5 +1,10 @@
+import sys
+from pathlib import Path
+
 import pytest
 
+sys.path.append(str(Path(__file__).parent / "code_contextvars"))
+
 pytest_plugins = ["pytester"]
 
 

tests/conftest.py

@@ -0,0 +1,27 @@
+import contextvars_for_test as ctxvars
+import pytest
+
+
+@pytest.fixture(scope="session", autouse=True)
+async def a_set_conftest_fixture_session_autouse():
+    ctxvars.a.set("a_set_conftest_fixture_session_autouse")
+
+
+@pytest.fixture(scope="session", autouse=True)
+async def b_set_conftest_cm_session_autouse():
+    ctxvars.b.set("b_set_conftest_cm_session_autouse")
+    yield
+
+
+@pytest.fixture()
+async def c_set_conftest_fixture_test():
+    ctxvars.c.set("c_set_conftest_fixture_test")
+
+
+@pytest.fixture()
+async def c_set_conftest_cm_test():
+    token = ctxvars.c.set("c_set_conftest_cm_test")
+    try:
+        yield
+    finally:
+        ctxvars.c.reset(token)

tests/test_contextvars/__init__.py

@@ -0,0 +1,98 @@
+# coding: spec
+
+import contextvars_for_test as ctxvars
+import pytest
+
+
+@pytest.fixture(scope="module", autouse=True)
+def d_set_conftest_cm_module_autouse():
+    token = ctxvars.d.set("d_set_conftest_fixture_test")
+    try:
+        yield
+    finally:
+        ctxvars.d.reset(token)
+
+
+@pytest.fixture(scope="module", autouse=True)
+async def e_set_conftest_cm_test():
+    assert ctxvars.f.get(ctxvars.Empty) is ctxvars.Empty
+    ctxvars.e.set("e_set_conftest_cm_module")
+    yield
+    assert ctxvars.f.get() == "f_set_conftest_cm_module"
+
+
+@pytest.fixture(scope="module", autouse=True)
+async def f_set_conftest_cm_module(e_set_conftest_cm_test):
+    assert ctxvars.e.get() == "e_set_conftest_cm_module"
+    ctxvars.f.set("f_set_conftest_cm_module")
+    yield
+
+
+@pytest.mark.order(1)
+async it "gets session modified vars":
+    assert ctxvars.a.get() == "a_set_conftest_fixture_session_autouse"
+    assert ctxvars.b.get() == "b_set_conftest_cm_session_autouse"
+    assert ctxvars.d.get() == "d_set_conftest_fixture_test"
+    assert ctxvars.e.get() == "e_set_conftest_cm_module"
+    assert ctxvars.f.get() == "f_set_conftest_cm_module"
+    ctxvars.assertVarsEmpty(excluding=("a", "b", "d", "e", "f"))
+
+
+@pytest.mark.order(2)
+async it "can use a fixture to change the var", c_set_conftest_fixture_test:
+    assert ctxvars.a.get() == "a_set_conftest_fixture_session_autouse"
+    assert ctxvars.b.get() == "b_set_conftest_cm_session_autouse"
+    assert ctxvars.d.get() == "d_set_conftest_fixture_test"
+    assert ctxvars.e.get() == "e_set_conftest_cm_module"
+    assert ctxvars.f.get() == "f_set_conftest_cm_module"
+
+    assert ctxvars.c.get() == "c_set_conftest_fixture_test"
+    ctxvars.assertVarsEmpty(excluding=("a", "b", "c", "d", "e", "f"))
+
+
+@pytest.mark.order(3)
+async it "does not reset contextvars for you":
+    """
+    It's too hard to know when a contextvar should be reset. It should
+    be up to whatever sets the contextvar to know when it should be unset
+    """
+    assert ctxvars.a.get() == "a_set_conftest_fixture_session_autouse"
+    assert ctxvars.b.get() == "b_set_conftest_cm_session_autouse"
+    assert ctxvars.d.get() == "d_set_conftest_fixture_test"
+    assert ctxvars.e.get() == "e_set_conftest_cm_module"
+    assert ctxvars.f.get() == "f_set_conftest_cm_module"
+
+    assert ctxvars.c.get() == "c_set_conftest_fixture_test"
+    ctxvars.assertVarsEmpty(excluding=("a", "b", "c", "d", "e", "f"))
+
+
+@pytest.mark.order(4)
+async it "works in context manager fixtures", c_set_conftest_cm_test:
+    """
+    It's too hard to know when a contextvar should be reset. It should
+    be up to whatever sets the contextvar to know when it should be unset
+    """
+    assert ctxvars.a.get() == "a_set_conftest_fixture_session_autouse"
+    assert ctxvars.b.get() == "b_set_conftest_cm_session_autouse"
+    assert ctxvars.d.get() == "d_set_conftest_fixture_test"
+    assert ctxvars.e.get() == "e_set_conftest_cm_module"
+    assert ctxvars.f.get() == "f_set_conftest_cm_module"
+
+    assert ctxvars.c.get() == "c_set_conftest_cm_test"
+    ctxvars.assertVarsEmpty(excluding=("a", "b", "c", "d", "e", "f"))
+
+
+@pytest.mark.order(5)
+it "resets the contextvar successfully when cm attempts that":
+    """
+    It's too hard to know when a contextvar should be reset. It should
+    be up to whatever sets the contextvar to know when it should be unset
+    """
+    assert ctxvars.a.get() == "a_set_conftest_fixture_session_autouse"
+    assert ctxvars.b.get() == "b_set_conftest_cm_session_autouse"
+    assert ctxvars.d.get() == "d_set_conftest_fixture_test"
+    assert ctxvars.e.get() == "e_set_conftest_cm_module"
+    assert ctxvars.f.get() == "f_set_conftest_cm_module"
+
+    assert ctxvars.c.get() == "c_set_conftest_fixture_test"
+    ctxvars.assertVarsEmpty(excluding=("a", "b", "c", "d", "e", "f"))

tests/test_contextvars/conftest.py

@@ -0,0 +1,14 @@
+# coding: spec
+
+import contextvars_for_test as ctxvars
+import pytest
+
+
+@pytest.mark.order(100)
+it "only keeps values from other module fixtures if they haven't been cleaned up":
+    assert ctxvars.a.get() == "a_set_conftest_fixture_session_autouse"
+    assert ctxvars.b.get() == "b_set_conftest_cm_session_autouse"
+    assert ctxvars.c.get() == "c_set_conftest_fixture_test"
+    assert ctxvars.e.get() == "e_set_conftest_cm_module"
+    assert ctxvars.f.get() == "f_set_conftest_cm_module"
+    ctxvars.assertVarsEmpty(excluding=("a", "b", "c", "e", "f"))