Restructure pytest plugin hooks

pytest_twisted.py

@@ -49,6 +49,11 @@ class _instances:
     reactor = None
 
 
+class _tracking:
+    async_yield_fixture_cache = {}
+    to_be_torn_down = []
+
+
 def _deprecate(deprecated, recommended):
     def decorator(f):
         @functools.wraps(f)
@@ -130,12 +135,6 @@ def stop_twisted_greenlet():
         _instances.gr_twisted.switch()
 
 
-class _CoroutineWrapper:
-    def __init__(self, coroutine, mark):
-        self.coroutine = coroutine
-        self.mark = mark
-
-
 def _marked_async_fixture(mark):
     @functools.wraps(pytest.fixture)
     def fixture(*args, **kwargs):
@@ -144,21 +143,14 @@ def fixture(*args, **kwargs):
         except IndexError:
             scope = kwargs.get('scope', 'function')
 
-        if scope != 'function':
+        if scope not in ['function', 'module']:
+            # TODO: add test for session scope (and that's it, right?)
+            #       then remove this and update docs
             raise AsyncFixtureUnsupportedScopeError.from_scope(scope=scope)
 
-        def marker(f):
-            @functools.wraps(f)
-            def w(*args, **kwargs):
-                return _CoroutineWrapper(
-                    coroutine=f(*args, **kwargs),
-                    mark=mark,
-                )
-
-            return w
-
         def decorator(f):
-            result = pytest.fixture(*args, **kwargs)(marker(f))
+            setattr(f, _mark_attribute_name, mark)
+            result = pytest.fixture(*args, **kwargs)(f)
 
             return result
 
@@ -167,61 +159,96 @@ def decorator(f):
     return fixture
 
 
+_mark_attribute_name = '_pytest_twisted_mark'
 async_fixture = _marked_async_fixture('async_fixture')
 async_yield_fixture = _marked_async_fixture('async_yield_fixture')
 
 
+def pytest_fixture_setup(fixturedef, request):
+    """Interface pytest to async setup for async and async yield fixtures."""
+    # TODO: what about inlineCallbacks fixtures?
+    maybe_mark = getattr(fixturedef.func, _mark_attribute_name, None)
+    if maybe_mark is None:
+        return None
+
+    mark = maybe_mark
+
+    run_inline_callbacks(_async_pytest_fixture_setup, fixturedef, request, mark)
+
+    return True
+
+
 @defer.inlineCallbacks
-def _pytest_pyfunc_call(pyfuncitem):
-    testfunction = pyfuncitem.obj
-    async_generators = []
-    funcargs = pyfuncitem.funcargs
-    if hasattr(pyfuncitem, "_fixtureinfo"):
-        testargs = {}
-        for arg in pyfuncitem._fixtureinfo.argnames:
-            if isinstance(funcargs[arg], _CoroutineWrapper):
-                wrapper = funcargs[arg]
-
-                if wrapper.mark == 'async_fixture':
-                    arg_value = yield defer.ensureDeferred(
-                        wrapper.coroutine
-                    )
-                elif wrapper.mark == 'async_yield_fixture':
-                    async_generators.append((arg, wrapper))
-                    arg_value = yield defer.ensureDeferred(
-                        wrapper.coroutine.__anext__(),
-                    )
-                else:
-                    raise UnrecognizedCoroutineMarkError.from_mark(
-                        mark=wrapper.mark,
-                    )
-            else:
-                arg_value = funcargs[arg]
-
-            testargs[arg] = arg_value
+def _async_pytest_fixture_setup(fixturedef, request, mark):
+    """Setup async and async yield fixtures."""
+    fixture_function = fixturedef.func
+
+    kwargs = {
+        name: request.getfixturevalue(name)
+        for name in fixturedef.argnames
+    }
+
+    if mark == 'async_fixture':
+        arg_value = yield defer.ensureDeferred(
+            fixture_function(**kwargs)
+        )
+    elif mark == 'async_yield_fixture':
+        coroutine = fixture_function(**kwargs)
+        # TODO: use request.addfinalizer() instead?
+        _tracking.async_yield_fixture_cache[request.param_index] = coroutine
+        arg_value = yield defer.ensureDeferred(
+            coroutine.__anext__(),
+        )
     else:
-        testargs = funcargs
-    result = yield testfunction(**testargs)
+        raise UnrecognizedCoroutineMarkError.from_mark(mark=mark)
 
-    async_generator_deferreds = [
-        (arg, defer.ensureDeferred(g.coroutine.__anext__()))
-        for arg, g in reversed(async_generators)
-    ]
+    fixturedef.cached_result = (arg_value, request.param_index, None)
 
-    for arg, d in async_generator_deferreds:
-        try:
-            yield d
-        except StopAsyncIteration:
-            continue
-        else:
-            raise AsyncGeneratorFixtureDidNotStopError.from_generator(
-                generator=arg,
-            )
+    defer.returnValue(arg_value)
 
-    defer.returnValue(result)
 
+# TODO: but don't we want to do the finalizer?  not wait until post it?
+def pytest_fixture_post_finalizer(fixturedef, request):
+    """Collect async yield fixture teardown requests for later handling."""
+    maybe_coroutine = _tracking.async_yield_fixture_cache.pop(
+        request.param_index,
+        None,
+    )
 
-def pytest_pyfunc_call(pyfuncitem):
+    if maybe_coroutine is None:
+        return None
+
+    coroutine = maybe_coroutine
+
+    deferred = defer.ensureDeferred(coroutine.__anext__())
+    _tracking.to_be_torn_down.append(deferred)
+    return None
+
+
+@defer.inlineCallbacks
+def tear_it_down(deferred):
+    """Tear down a specific async yield fixture."""
+    try:
+        yield deferred
+    except StopAsyncIteration:
+        return
+    except Exception as e:
+        e = e
+    else:
+        e = None
+
+    # TODO: six.raise_from()
+    raise AsyncGeneratorFixtureDidNotStopError.from_generator(
+        generator=deferred,
+    )
+
+
+# TODO: https://docs.pytest.org/en/latest/reference.html#_pytest.hookspec.pytest_runtest_protocol
+#       claims it should also take a nextItem but that triggers a direct error
+
+
+def run_inline_callbacks(f, *args):
+    """Interface into Twisted greenlet to run and wait for deferred."""
     if _instances.gr_twisted is not None:
         if _instances.gr_twisted.dead:
             raise RuntimeError("twisted reactor has stopped")
@@ -230,26 +257,50 @@ def in_reactor(d, f, *args):
             return defer.maybeDeferred(f, *args).chainDeferred(d)
 
         d = defer.Deferred()
-        _instances.reactor.callLater(
-            0.0, in_reactor, d, _pytest_pyfunc_call, pyfuncitem
-        )
+        _instances.reactor.callLater(0.0, in_reactor, d, f, *args)
         blockon_default(d)
     else:
         if not _instances.reactor.running:
             raise RuntimeError("twisted reactor is not running")
-        blockingCallFromThread(
-            _instances.reactor, _pytest_pyfunc_call, pyfuncitem
-        )
+        blockingCallFromThread(_instances.reactor, f, *args)
+
+
+@pytest.hookimpl(hookwrapper=True)
+def pytest_runtest_teardown(item):
+    """Tear down collected async yield fixtures."""
+    yield
+
+    while len(_tracking.to_be_torn_down) > 0:
+        deferred = _tracking.to_be_torn_down.pop(0)
+        run_inline_callbacks(tear_it_down, deferred)
+
+
+def pytest_pyfunc_call(pyfuncitem):
+    """Interface to async test call handler."""
+    # TODO: only handle 'our' tests?  what is the point of handling others?
+    run_inline_callbacks(_async_pytest_pyfunc_call, pyfuncitem)
     return True
 
 
+@defer.inlineCallbacks
+def _async_pytest_pyfunc_call(pyfuncitem):
+    """Run test function."""
+    kwargs = {
+        name: value
+        for name, value in pyfuncitem.funcargs.items()
+        if name in pyfuncitem._fixtureinfo.argnames
+    }
+    result = yield pyfuncitem.obj(**kwargs)
+    defer.returnValue(result)
+
+
 @pytest.fixture(scope="session", autouse=True)
-def twisted_greenlet(request):
-    request.addfinalizer(stop_twisted_greenlet)
+def twisted_greenlet():
     return _instances.gr_twisted
 
 
 def init_default_reactor():
+    """Install the default Twisted reactor."""
     import twisted.internet.default
 
     module = inspect.getmodule(twisted.internet.default.install)
@@ -265,6 +316,7 @@ def init_default_reactor():
 
 
 def init_qt5_reactor():
+    """Install the qt5reactor...  reactor."""
     import qt5reactor
 
     _install_reactor(
@@ -273,6 +325,7 @@ def init_qt5_reactor():
 
 
 def init_asyncio_reactor():
+    """Install the Twisted reactor for asyncio."""
     from twisted.internet import asyncioreactor
 
     _install_reactor(
@@ -289,6 +342,7 @@ def init_asyncio_reactor():
 
 
 def _install_reactor(reactor_installer, reactor_type):
+    """Install the specified reactor and create the greenlet."""
     try:
         reactor_installer()
     except error.ReactorAlreadyInstalledError:
@@ -308,6 +362,7 @@ def _install_reactor(reactor_installer, reactor_type):
 
 
 def pytest_addoption(parser):
+    """Add options into the pytest CLI."""
     group = parser.getgroup("twisted")
     group.addoption(
         "--reactor",
@@ -317,6 +372,7 @@ def pytest_addoption(parser):
 
 
 def pytest_configure(config):
+    """Identify and install chosen reactor."""
     pytest.inlineCallbacks = _deprecate(
         deprecated='pytest.inlineCallbacks',
         recommended='pytest_twisted.inlineCallbacks',
@@ -329,7 +385,13 @@ def pytest_configure(config):
     reactor_installers[config.getoption("reactor")]()
 
 
+def pytest_unconfigure(config):
+    """Stop the reactor greenlet."""
+    stop_twisted_greenlet()
+
+
 def _use_asyncio_selector_if_required(config):
+    """Set asyncio selector event loop policy if needed."""
     # https://twistedmatrix.com/trac/ticket/9766
     # https://github.com/pytest-dev/pytest-twisted/issues/80