How to timeout an async test in pytest with fixture?

Question

I am testing an async function that might get deadlocked. I tried to add a fixture to limit the function to only run for 5 seconds before raising a failure, but it hasn't worked so far.

Setup:

pipenv --python==3.6
pipenv install pytest==4.4.1
pipenv install pytest-asyncio==0.10.0

Code:

import asyncio
import pytest

@pytest.fixture
def my_fixture():
  # attempt to start a timer that will stop the test somehow
  asyncio.ensure_future(time_limit())
  yield 'eggs'


async def time_limit():
  await asyncio.sleep(5)
  print('time limit reached')     # this isn't printed
  raise AssertionError


@pytest.mark.asyncio
async def test(my_fixture):
  assert my_fixture == 'eggs'
  await asyncio.sleep(10)
  print('this should not print')  # this is printed
  assert 0

--

Edit: Mikhail's solution works fine. I can't find a way to incorporate it into a fixture, though.

Answer 1

Here is my hackish solution

It's using a pytest marker to mark tests for timeout, and an automatic fixture that will only apply if the mark is present

It will run a separate task that will cancel the test task after a delay

async def cancel_test_after_timeout(test_name: str, timeout: float):
    """
    Ran into a separate task: after a given time, cancel a running test coroutine
    with an error message

    The "cancel" task itself will be canceled after the test is finished
    """
    await asyncio.sleep(timeout)

    for task in asyncio.all_tasks():
        coro_name: str = task.get_coro().__name__  # type: ignore

        if coro_name == test_name:
            task.cancel(f"Canceling {coro_name} after {timeout:.2g}s")


@pytest.fixture(autouse=True)
async def async_timeout_on_mark(request):
    """
    In case test has the async_timeout marker, create a new task that will cancel
    the running test after the timeout given as parameter

    """
    marker = request.node.get_closest_marker("async_timeout")
    if not marker:
        yield
        return

    test_name = request.function.__name__
    timeout = marker.args[0] if marker.args else 30.0

    task = asyncio.create_task(cancel_test_after_timeout(test_name, timeout))
    yield
    task.cancel()


def pytest_configure(config: pytest.Config):
    config.addinivalue_line(
        "markers",
        "async_timeout(timeout): cancels the test execution after the"
        "specified amount of seconds",
    )

A simple example:

@pytest.mark.parametrize("wait_time", (0.1, 1.1))
@pytest.mark.asyncio()
@pytest.mark.async_timeout(1)
async def test_timeout_marker(wait_time):
    # the second test will fail
    print(f"waiting {wait_time}")
    await asyncio.sleep(wait_time)

Answer 2

Instead of using a fixture, I solved it this way using a decorator:

def timeout(delay):
    def decorator(func):
        @wraps(func)
        async def new_func(*args, **kwargs):
            async with asyncio.timeout(delay):
                return await func(*args, **kwargs)

        return new_func

    return decorator
    

@pytest.mark.asyncio
@timeout(3)
async def test_forever_fails():
    await asyncio.Future()

Requires python 3.11

Or I believe trio provides something like this for earlier python versions.

Answer 3

I just loved Quimby's approach of marking tests with timeouts. Here's my attempt to improve it, using pytest marks:

# tests/conftest.py
import asyncio


@pytest.hookimpl(tryfirst=True, hookwrapper=True)
def pytest_pyfunc_call(pyfuncitem: pytest.Function):
    """
    Wrap all tests marked with pytest.mark.async_timeout with their specified timeout.
    """
    orig_obj = pyfuncitem.obj

    if marker := pyfuncitem.get_closest_marker("async_timeout"):

        async def new_obj(*args, **kwargs):
            """Timed wrapper around the test function."""
            try:
                return await asyncio.wait_for(orig_obj(*args, **kwargs), timeout=marker.args[0])
            except (asyncio.CancelledError, asyncio.TimeoutError):
                pytest.fail(f"Test {pyfuncitem.name} did not finish in time.")

        pyfuncitem.obj = new_obj

    yield


def pytest_configure(config: pytest.Config):
    config.addinivalue_line("markers", "async_timeout(timeout): cancels the test execution after the specified amount of seconds")

Usage:

@pytest.mark.asyncio
@pytest.mark.async_timeout(10)
async def potentially_hanging_function():
    await asyncio.sleep(20)

It should not be hard to include this to the asyncio mark on pytest-asyncio, so we can get a syntax like:

@pytest.mark.asyncio(timeout=10)
async def potentially_hanging_function():
    await asyncio.sleep(20)

EDIT: looks like there's already a PR for that.

Answer 4

There is a way to use fixtures for timeout, one just needs to add the following hook into conftest.py.

• Any fixture prefixed with timeout must return a number of seconds(int, float) the test can run.
• The closest fixture w.r.t scope is chosen. autouse fixtures have lesser priority than explicitly chosen ones. Later one is preferred. Unfortunately order in the function argument list does NOT matter.
• If there is no such fixture, the test is not restricted and will run indefinitely as usual.
• The test must be marked with pytest.mark.asyncio too, but that is needed anyway.

# Add to conftest.py
import asyncio

import pytest

_TIMEOUT_FIXTURE_PREFIX = "timeout"


@pytest.hookimpl(tryfirst=True, hookwrapper=True)
def pytest_runtest_setup(item: pytest.Item):
    """Wrap all tests marked with pytest.mark.asyncio with their specified timeout.

    Must run as early as possible.

    Parameters
    ----------
    item : pytest.Item
        Test to wrap
    """
    yield
    orig_obj = item.obj
    timeouts = [n for n in item.funcargs if n.startswith(_TIMEOUT_FIXTURE_PREFIX)]
    # Picks the closest timeout fixture if there are multiple
    tname = None if len(timeouts) == 0 else timeouts[-1]

    # Only pick marked functions
    if item.get_closest_marker("asyncio") is not None and tname is not None:

        async def new_obj(*args, **kwargs):
            """Timed wrapper around the test function."""
            try:
                return await asyncio.wait_for(
                    orig_obj(*args, **kwargs), timeout=item.funcargs[tname]
                )
            except Exception as e:
                pytest.fail(f"Test {item.name} did not finish in time.")

        item.obj = new_obj

Example:

@pytest.fixture
def timeout_2s():
    return 2


@pytest.fixture(scope="module", autouse=True)
def timeout_5s():
    # You can do whatever you need here, just return/yield a number
    return 5


async def test_timeout_1():
    # Uses timeout_5s fixture by default
    await aio.sleep(0)  # Passes
    return 1


async def test_timeout_2(timeout_2s):
    # Uses timeout_2s because it is closest
    await aio.sleep(5)  # Timeouts

WARNING

Might not work with some other plugins, I have only tested it with pytest-asyncio, it definitely won't work if item is redefined by some hook.

Answer 5

Convenient way to limit function (or block of code) with timeout is to use async-timeout module. You can use it inside your test function or, for example, create a decorator. Unlike with fixture it'll allow to specify concrete time for each test:

import asyncio
import pytest
from async_timeout import timeout


def with_timeout(t):
    def wrapper(corofunc):
        async def run(*args, **kwargs):
            with timeout(t):
                return await corofunc(*args, **kwargs)
        return run       
    return wrapper


@pytest.mark.asyncio
@with_timeout(2)
async def test_sleep_1():
    await asyncio.sleep(1)
    assert 1 == 1


@pytest.mark.asyncio
@with_timeout(2)
async def test_sleep_3():
    await asyncio.sleep(3)
    assert 1 == 1

It's not hard to create decorator for concrete time (with_timeout_5 = partial(with_timeout, 5)).

---

I don't know how to create texture (if you really need fixture), but code above can provide starting point. Also not sure if there's a common way to achieve goal better.