Merge pull request #28 from touilleMan/fixture-in-recursive-acontextmanager-with-nursery-fixture

Recursive async context managers for fixtures + nursery fixture

Author: touilleMan

pytest_trio/_tests/test_async_yield_fixture.py

@@ -224,3 +224,40 @@ async def test_actual_test(fix1):
     # TODO: should trigger error instead of failure
     # result.assert_outcomes(error=1)
     result.assert_outcomes(failed=1)
+
+
+@pytest.mark.skipif(sys.version_info < (3, 6), reason="requires python3.6")
+def test_async_yield_fixture_with_nursery(testdir):
+
+    testdir.makepyfile(
+        """
+        import pytest
+        import trio
+
+
+        async def handle_client(stream):
+            while True:
+                buff = await stream.receive_some(4)
+                await stream.send_all(buff)
+
+
+        @pytest.fixture
+        async def server():
+            async with trio.open_nursery() as nursery:
+                listeners = await nursery.start(trio.serve_tcp, handle_client, 0)
+                yield listeners[0]
+                nursery.cancel_scope.cancel()
+
+
+        @pytest.mark.trio
+        async def test_actual_test(server):
+            stream = await trio.testing.open_stream_to_socket_listener(server)
+            await stream.send_all(b'ping')
+            rep = await stream.receive_some(4)
+            assert rep == b'ping'
+    """
+    )
+
+    result = testdir.runpytest()
+
+    result.assert_outcomes(passed=1)

pytest_trio/_tests/test_nursery_fixture.py

@@ -0,0 +1,22 @@
+import pytest
+import trio
+
+
+async def handle_client(stream):
+    while True:
+        buff = await stream.receive_some(4)
+        await stream.send_all(buff)
+
+
+@pytest.fixture
+async def server(nursery):
+    listeners = await nursery.start(trio.serve_tcp, handle_client, 0)
+    return listeners[0]
+
+
+@pytest.mark.trio
+async def test_try(server):
+    stream = await trio.testing.open_stream_to_socket_listener(server)
+    await stream.send_all(b'ping')
+    rep = await stream.receive_some(4)
+    assert rep == b'ping'

pytest_trio/_tests/test_sync_fixture.py

@@ -44,3 +44,46 @@ def test_after():
     result = testdir.runpytest()
 
     result.assert_outcomes(passed=3)
+
+
+def test_single_yield_fixture_with_async_deps(testdir):
+
+    testdir.makepyfile(
+        """
+        import pytest
+        import trio
+
+        events = []
+
+        @pytest.fixture
+        async def fix0():
+            events.append('fix0 setup')
+            await trio.sleep(0)
+            return 'fix0'
+
+        @pytest.fixture
+        def fix1(fix0):
+            events.append('fix1 setup')
+            yield 'fix1 - ' + fix0
+            events.append('fix1 teardown')
+
+        def test_before():
+            assert not events
+
+        @pytest.mark.trio
+        async def test_actual_test(fix1):
+            assert events == ['fix0 setup', 'fix1 setup']
+            assert fix1 == 'fix1 - fix0'
+
+        def test_after():
+            assert events == [
+                'fix0 setup',
+                'fix1 setup',
+                'fix1 teardown',
+            ]
+    """
+    )
+
+    result = testdir.runpytest()
+
+    result.assert_outcomes(passed=3)

pytest_trio/plugin.py

@@ -11,7 +11,9 @@ def isasyncgenfunction(x):
 
 import pytest
 import trio
+from trio._util import acontextmanager
 from trio.testing import MockClock, trio_test
+from async_generator import async_generator, yield_
 
 
 def pytest_configure(config):
@@ -29,34 +31,69 @@ def _trio_test_runner_factory(item):
     @trio_test
     async def _bootstrap_fixture_and_run_test(**kwargs):
         __tracebackhide__ = True
-        resolved_kwargs = await _setup_async_fixtures_in(kwargs)
-        await testfunc(**resolved_kwargs)
-        await _teardown_async_fixtures_in(kwargs)
+        user_exc = None
+        # Open the nursery exposed as fixture
+        async with trio.open_nursery() as nursery:
+            item._trio_nursery = nursery
+            try:
+                async with _setup_async_fixtures_in(kwargs) as resolved_kwargs:
+                    try:
+                        await testfunc(**resolved_kwargs)
+                    except BaseException as exc:
+                        # Regular pytest fixture don't have access to the test
+                        # exception in there teardown, we mimic this behavior here.
+                        user_exc = exc
+            except BaseException as exc:
+                # If we are here, the exception comes from the fixtures setup
+                # or teardown
+                if user_exc:
+                    raise exc from user_exc
+                else:
+                    raise exc
+            finally:
+                # No matter what the nursery fixture should be closed when test is over
+                nursery.cancel_scope.cancel()
+
+        # Finally re-raise or original exception coming from the test if needed
+        if user_exc:
+            raise user_exc
 
     return _bootstrap_fixture_and_run_test
 
 
+@acontextmanager
+@async_generator
 async def _setup_async_fixtures_in(deps):
-    resolved_deps = {**deps}
+    __tracebackhide__ = True
 
-    async def _resolve_and_update_deps(afunc, deps, entry):
-        deps[entry] = await afunc()
+    need_resolved_deps_stack = [
+        (k, v) for k, v in deps.items() if isinstance(v, BaseAsyncFixture)
+    ]
 
-    async with trio.open_nursery() as nursery:
-        for depname, depval in resolved_deps.items():
-            if isinstance(depval, BaseAsyncFixture):
-                nursery.start_soon(
-                    _resolve_and_update_deps, depval.setup, resolved_deps,
-                    depname
-                )
-    return resolved_deps
+    if not need_resolved_deps_stack:
+        await yield_(deps)
+        return
 
-
-async def _teardown_async_fixtures_in(deps):
-    async with trio.open_nursery() as nursery:
-        for depval in deps.values():
-            if isinstance(depval, BaseAsyncFixture):
-                nursery.start_soon(depval.teardown)
+    @acontextmanager
+    @async_generator
+    async def _recursive_setup(deps_stack):
+        __tracebackhide__ = True
+        name, dep = deps_stack.pop()
+        async with dep.setup() as resolved:
+            if not deps_stack:
+                await yield_([(name, resolved)])
+            else:
+                async with _recursive_setup(
+                    deps_stack
+                ) as remains_deps_stack_resolved:
+                    await yield_(
+                        remains_deps_stack_resolved + [(name, resolved)]
+                    )
+
+    async with _recursive_setup(
+        need_resolved_deps_stack
+    ) as resolved_deps_stack:
+        await yield_({**deps, **dict(resolved_deps_stack)})
 
 
 class BaseAsyncFixture:
@@ -68,49 +105,41 @@ def __init__(self, fixturedef, deps={}):
         self.fixturedef = fixturedef
         self.deps = deps
         self.setup_done = False
-        self.teardown_done = False
         self.result = None
-        self.lock = trio.Lock()
 
+    @acontextmanager
+    @async_generator
     async def setup(self):
-        async with self.lock:
-            if not self.setup_done:
-                self.result = await self._setup()
-                self.setup_done = True
-            return self.result
+        __tracebackhide__ = True
+        if self.setup_done:
+            await yield_(self.result)
+        else:
+            async with _setup_async_fixtures_in(self.deps) as resolved_deps:
+                async with self._setup(resolved_deps) as self.result:
+                    self.setup_done = True
+                    await yield_(self.result)
 
     async def _setup(self):
         raise NotImplementedError()
 
-    async def teardown(self):
-        async with self.lock:
-            if not self.teardown_done:
-                await self._teardown()
-                self.teardown_done = True
-
-    async def _teardown(self):
-        raise NotImplementedError()
-
 
 class AsyncYieldFixture(BaseAsyncFixture):
     """
     Async generator fixture.
     """
 
-    def __init__(self, *args):
-        super().__init__(*args)
-        self.agen = None
+    @acontextmanager
+    @async_generator
+    async def _setup(self, resolved_deps):
+        __tracebackhide__ = True
+        agen = self.fixturedef.func(**resolved_deps)
 
-    async def _setup(self):
-        resolved_deps = await _setup_async_fixtures_in(self.deps)
-        self.agen = self.fixturedef.func(**resolved_deps)
-        return await self.agen.asend(None)
+        await yield_(await agen.asend(None))
 
-    async def _teardown(self):
         try:
-            await self.agen.asend(None)
+            await agen.asend(None)
         except StopAsyncIteration:
-            await _teardown_async_fixtures_in(self.deps)
+            pass
         else:
             raise RuntimeError('Only one yield in fixture is allowed')
 
@@ -120,33 +149,30 @@ class SyncFixtureWithAsyncDeps(BaseAsyncFixture):
     Synchronous function fixture with asynchronous dependencies fixtures.
     """
 
-    async def _setup(self):
-        resolved_deps = await _setup_async_fixtures_in(self.deps)
-        return self.fixturedef.func(**resolved_deps)
-
-    async def _teardown(self):
-        await _teardown_async_fixtures_in(self.deps)
+    @acontextmanager
+    @async_generator
+    async def _setup(self, resolved_deps):
+        __tracebackhide__ = True
+        await yield_(self.fixturedef.func(**resolved_deps))
 
 
 class SyncYieldFixtureWithAsyncDeps(BaseAsyncFixture):
     """
     Synchronous generator fixture with asynchronous dependencies fixtures.
     """
 
-    def __init__(self, *args):
-        super().__init__(*args)
-        self.agen = None
+    @acontextmanager
+    @async_generator
+    async def _setup(self, resolved_deps):
+        __tracebackhide__ = True
+        gen = self.fixturedef.func(**resolved_deps)
 
-    async def _setup(self):
-        resolved_deps = await _setup_async_fixtures_in(self.deps)
-        self.gen = self.fixturedef.func(**resolved_deps)
-        return self.gen.send(None)
+        await yield_(gen.send(None))
 
-    async def _teardown(self):
         try:
-            await self.gen.send(None)
+            gen.send(None)
         except StopIteration:
-            await _teardown_async_fixtures_in(self.deps)
+            pass
         else:
             raise RuntimeError('Only one yield in fixture is allowed')
 
@@ -156,12 +182,11 @@ class AsyncFixture(BaseAsyncFixture):
     Regular async fixture (i.e. coroutine).
     """
 
-    async def _setup(self):
-        resolved_deps = await _setup_async_fixtures_in(self.deps)
-        return await self.fixturedef.func(**resolved_deps)
-
-    async def _teardown(self):
-        await _teardown_async_fixtures_in(self.deps)
+    @acontextmanager
+    @async_generator
+    async def _setup(self, resolved_deps):
+        __tracebackhide__ = True
+        await yield_(await self.fixturedef.func(**resolved_deps))
 
 
 def _install_async_fixture_if_needed(fixturedef, request):
@@ -215,3 +240,8 @@ def mock_clock():
 @pytest.fixture
 def autojump_clock():
     return MockClock(autojump_threshold=0)
+
+
+@pytest.fixture
+async def nursery(request):
+    return request.node._trio_nursery

setup.py

@@ -14,9 +14,10 @@
     author_email="[email protected]",
     license="MIT -or- Apache License 2.0",
     packages=find_packages(),
-    entry_points={'pytest11': ['trio = pytest_trio.plugin',]},
+    entry_points={'pytest11': ['trio = pytest_trio.plugin']},
     install_requires=[
         "trio",
+        "async_generator >= 1.6",
     ],
     keywords=[
         'async',