plugin.py

"""pytest-asyncio implementation."""
import asyncio
import contextlib
import functools
import inspect
import socket

import pytest

try:
    from _pytest.python import transfer_markers
except ImportError:  # Pytest 4.1.0 removes the transfer_marker api (#104)

    def transfer_markers(*args, **kwargs):  # noqa
        """Noop when over pytest 4.1.0"""
        pass


from inspect import isasyncgenfunction


def _is_coroutine(obj):
    """Check to see if an object is really an asyncio coroutine."""
    return asyncio.iscoroutinefunction(obj) or inspect.isgeneratorfunction(obj)


def pytest_configure(config):
    """Inject documentation."""
    config.addinivalue_line(
        "markers",
        "asyncio: "
        "mark the test as a coroutine, it will be "
        "run using an asyncio event loop",
    )


@pytest.mark.tryfirst
def pytest_pycollect_makeitem(collector, name, obj):
    """A pytest hook to collect asyncio coroutines."""
    if collector.funcnamefilter(name) and _is_coroutine(obj):
        item = pytest.Function.from_parent(collector, name=name)

        # Due to how pytest test collection works, module-level pytestmarks
        # are applied after the collection step. Since this is the collection
        # step, we look ourselves.
        transfer_markers(obj, item.cls, item.module)
        item = pytest.Function.from_parent(collector, name=name)  # To reload keywords.

        if "asyncio" in item.keywords:
            return list(collector._genfunctions(name, obj))


class FixtureStripper:
    """Include additional Fixture, and then strip them"""

    REQUEST = "request"
    EVENT_LOOP = "event_loop"

    def __init__(self, fixturedef):
        self.fixturedef = fixturedef
        self.to_strip = set()

    def add(self, name):
        """Add fixture name to fixturedef
        and record in to_strip list (If not previously included)"""
        if name in self.fixturedef.argnames:
            return
        self.fixturedef.argnames += (name,)
        self.to_strip.add(name)

    def get_and_strip_from(self, name, data_dict):
        """Strip name from data, and return value"""
        result = data_dict[name]
        if name in self.to_strip:
            del data_dict[name]
        return result


@pytest.hookimpl(trylast=True)
def pytest_fixture_post_finalizer(fixturedef, request):
    """Called after fixture teardown"""
    if fixturedef.argname == "event_loop":
        # Set empty loop policy, so that subsequent get_event_loop() provides a new loop
        asyncio.set_event_loop_policy(None)


@pytest.hookimpl(hookwrapper=True)
def pytest_fixture_setup(fixturedef, request):
    """Adjust the event loop policy when an event loop is produced."""
    if fixturedef.argname == "event_loop":
        outcome = yield
        loop = outcome.get_result()
        policy = asyncio.get_event_loop_policy()
        try:
            old_loop = policy.get_event_loop()
            if old_loop is not loop:
                old_loop.close()
        except RuntimeError:
            # Swallow this, since it's probably bad event loop hygiene.
            pass
        policy.set_event_loop(loop)
        return

    if isasyncgenfunction(fixturedef.func):
        # This is an async generator function. Wrap it accordingly.
        generator = fixturedef.func

        fixture_stripper = FixtureStripper(fixturedef)
        fixture_stripper.add(FixtureStripper.EVENT_LOOP)
        fixture_stripper.add(FixtureStripper.REQUEST)

        def wrapper(*args, **kwargs):
            loop = fixture_stripper.get_and_strip_from(
                FixtureStripper.EVENT_LOOP, kwargs
            )
            request = fixture_stripper.get_and_strip_from(
                FixtureStripper.REQUEST, kwargs
            )

            gen_obj = generator(*args, **kwargs)

            async def setup():
                res = await gen_obj.__anext__()
                return res

            def finalizer():
                """Yield again, to finalize."""

                async def async_finalizer():
                    try:
                        await gen_obj.__anext__()
                    except StopAsyncIteration:
                        pass
                    else:
                        msg = "Async generator fixture didn't stop."
                        msg += "Yield only once."
                        raise ValueError(msg)

                loop.run_until_complete(async_finalizer())

            request.addfinalizer(finalizer)
            return loop.run_until_complete(setup())

        fixturedef.func = wrapper
    elif inspect.iscoroutinefunction(fixturedef.func):
        coro = fixturedef.func

        fixture_stripper = FixtureStripper(fixturedef)
        fixture_stripper.add(FixtureStripper.EVENT_LOOP)

        def wrapper(*args, **kwargs):
            loop = fixture_stripper.get_and_strip_from(
                FixtureStripper.EVENT_LOOP, kwargs
            )

            async def setup():
                res = await coro(*args, **kwargs)
                return res

            return loop.run_until_complete(setup())

        fixturedef.func = wrapper
    yield


@pytest.hookimpl(tryfirst=True, hookwrapper=True)
def pytest_pyfunc_call(pyfuncitem):
    """
    Run asyncio marked test functions in an event loop instead of a normal
    function call.
    """
    if "asyncio" in pyfuncitem.keywords:
        if getattr(pyfuncitem.obj, "is_hypothesis_test", False):
            pyfuncitem.obj.hypothesis.inner_test = wrap_in_sync(
                pyfuncitem.obj.hypothesis.inner_test,
                _loop=pyfuncitem.funcargs["event_loop"],
            )
        else:
            pyfuncitem.obj = wrap_in_sync(
                pyfuncitem.obj, _loop=pyfuncitem.funcargs["event_loop"]
            )
    yield


def wrap_in_sync(func, _loop):
    """Return a sync wrapper around an async function executing it in the
    current event loop."""

    @functools.wraps(func)
    def inner(**kwargs):
        coro = func(**kwargs)
        if coro is not None:
            task = asyncio.ensure_future(coro, loop=_loop)
            try:
                _loop.run_until_complete(task)
            except BaseException:
                # run_until_complete doesn't get the result from exceptions
                # that are not subclasses of `Exception`. Consume all
                # exceptions to prevent asyncio's warning from logging.
                if task.done() and not task.cancelled():
                    task.exception()
                raise

    return inner


def pytest_runtest_setup(item):
    if "asyncio" in item.keywords:
        # inject an event loop fixture for all async tests
        if "event_loop" in item.fixturenames:
            item.fixturenames.remove("event_loop")
        item.fixturenames.insert(0, "event_loop")
    if (
        item.get_closest_marker("asyncio") is not None
        and not getattr(item.obj, "hypothesis", False)
        and getattr(item.obj, "is_hypothesis_test", False)
    ):
        pytest.fail(
            "test function `%r` is using Hypothesis, but pytest-asyncio "
            "only works with Hypothesis 3.64.0 or later." % item
        )


@pytest.fixture
def event_loop(request):
    """Create an instance of the default event loop for each test case."""
    loop = asyncio.get_event_loop_policy().new_event_loop()
    yield loop
    loop.close()


def _unused_tcp_port():
    """Find an unused localhost TCP port from 1024-65535 and return it."""
    with contextlib.closing(socket.socket()) as sock:
        sock.bind(("127.0.0.1", 0))
        return sock.getsockname()[1]


@pytest.fixture
def unused_tcp_port():
    return _unused_tcp_port()


@pytest.fixture(scope="session")
def unused_tcp_port_factory():
    """A factory function, producing different unused TCP ports."""
    produced = set()

    def factory():
        """Return an unused port."""
        port = _unused_tcp_port()

        while port in produced:
            port = _unused_tcp_port()

        produced.add(port)

        return port

    return factory