Shy trio fixtures

test_shy.py

import pytest
import trio
from contextlib import contextmanager
from functools import wraps
from pytest_trio.plugin import TrioFixture
from inspect import isasyncgenfunction


def shy_in_presence_of(*names):
    def decorator(func):
        @wraps(func)
        async def wrapper(request, *args, **kwargs):
            for name in names:
                if name in request.fixturenames:
                    fixture = request.getfixturevalue(name)
                    if isinstance(fixture, TrioFixture):
                        await fixture.setup_done.wait()
            if isasyncgenfunction(func):
                async for value in func(request, *args, **kwargs):
                    yield value
            else:
                yield func(request, *args, **kwargs)
        return wrapper
    return decorator


finish_line = []

@contextmanager
def run_for_it(animal):
    finish_line.append(animal)
    yield
    finish_line.remove(animal)


@pytest.fixture
async def sloth():
    await trio.sleep(2)
    with run_for_it("sloth"):
        yield

@pytest.fixture
async def turtle():
    await trio.sleep(2)
    with run_for_it("turtle"):
        yield

@pytest.fixture
@shy_in_presence_of("turtle")
async def rabbit(request):
    await trio.sleep(1)
    with run_for_it("rabbit"):
        yield


async def test_rabbit_beats_sloth(rabbit, sloth):
    assert finish_line == ["rabbit", "sloth"]

async def test_rabbit_too_shy_to_beat_turtle(rabbit, turtle):
    assert finish_line == ["turtle", "rabbit"]