asyncio: #8 consider contextvars

Make it so everything gets executed in the same asyncio context

Author: delfick

alt_pytest_asyncio/async_converters.py

@@ -14,16 +14,75 @@
 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)
+        fixturedef.func.__alt_asyncio_pytest_converted__ = True
 
     elif inspect.isasyncgenfunction(fixturedef.func):
         convert_async_gen_fixture(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)
+
+    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(test_tasks, func, timeout, *args, **kwargs):

alt_pytest_asyncio/plugin.py

@@ -1,3 +1,12 @@
+from alt_pytest_asyncio.async_converters import convert_fixtures, converted_async_test
+from alt_pytest_asyncio import async_converters
+
+from _pytest._code.code import ExceptionInfo
+from functools import partial, wraps
+from collections import defaultdict
+from unittest import mock
+import contextvars
+import inspect
 import asyncio
 import inspect
 import sys
@@ -19,14 +28,48 @@ 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"""
         config.addinivalue_line(
             "markers", "async_timeout(length): mark async test to have a timeout"
         )
 
+    @pytest.hookimpl(tryfirst=True, hookwrapper=True)
+    def pytest_sessionstart(self, session):
+        """
+        Python won't let me create the task with it's own context And
+        run_until_complete doesn't let me propagate a context.
+
+        So I monkey patch call_soon to force the context So I don't need my own
+        Task class
+
+        This obviously won't work if Task doesn't something other than
+        ``call_soon(self._step)``.
+        """
+        original = self.loop.call_soon
+
+        @wraps(original)
+        def wrapped_call_soon(callback, *args, context=None):
+            is_alt_pytest = False
+            for frm in inspect.stack():
+                mod = inspect.getmodule(frm[0])
+                if mod is async_converters:
+                    is_alt_pytest = True
+                    break
+
+            if is_alt_pytest:
+                return original(callback, *args, context=self.ctx)
+            else:
+                return original(callback, *args, context=context)
+
+        self.loop_patch = mock.patch.object(self.loop, "call_soon", wrapped_call_soon)
+        self.loop_patch.start()
+        yield
+
     @pytest.hookimpl(tryfirst=True, hookwrapper=True)
     def pytest_sessionfinish(self, session, exitstatus):
         """
@@ -54,10 +97,13 @@ def pytest_sessionfinish(self, session, exitstatus):
                 finally:
                     self.loop.close()
 
+            if hasattr(self, "loop_patch"):
+                self.loop_patch.stop()
+
     @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)
@@ -76,6 +122,22 @@ def pytest_pyfunc_call(self, pyfuncitem):
 
             o = pyfuncitem.obj
             pyfuncitem.obj = wraps(o)(partial(converted_async_test, 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 @@
+from contextvars import ContextVar
+import string
+
+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,4 +1,12 @@
 import pytest
+from pathlib import Path
+import asyncio
+import pytest
+import struct
+import socket
+import sys
+
+sys.path.append(str(Path(__file__).parent / "code_contextvars"))
 
 pytest_plugins = ["pytester"]
 

tests/conftest.py

@@ -0,0 +1,28 @@
+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"))