Unit testing async functions that loop indefinitely

Help and advice
1

Hi there,

I’m a new user of this library and I’m trying to wrap my head around writing unit tests for async functions that loop indefinitely.

Background

For context, I have a function run_task that will be called multiple times by nursery.start_soon. There will be a saver function to periodically save some data until all run_task function calls have been awaited.

# test_saver.py
from typing import Callable, Awaitable
from pathlib import Path

import trio

async def run_task(sleep_time: int) -> None:
    await trio.sleep(sleep_time)


async def saver(nursery: trio.Nursery, target: Callable[..., Awaitable], savepath: Path) -> None:
    """Save data periodically until all `target` functions have finished executing."""
    while any(target.__name__ in task.name for task in nursery.child_tasks):
        with open(savepath, "w") as f:
            f.write(str(trio.current_time()))
        await trio.sleep(1)
    nursery.cancel_scope.cancel()

Problem

I want to test that saver cancels the nursery after 3 secs, which is when all run_task tasks have completed (in the examples below).

However, I am unable to achieve such granularity in my tests — I only can verify that saver cancels the nursery eventually, not exactly / around the elapsed time of 3 secs.

Here are some of my attempts to do so (including using pytest-trio fixtures to no avail)

# continuation of test_saver.py

async def test_basic(tmp_path):
    # PASSING: But how can I check that the nursery was cancelled at around 3 sec?
    async with trio.open_nursery() as nursery:
        for i in range(3):
            nursery.start_soon(run_task, i)
        nursery.start_soon(saver, nursery, run_task, tmp_path / 'save.txt')
    assert nursery.cancel_scope.cancel_called


async def test_with_nursery_fixture(tmp_path, nursery):
    # FAILING: Might be due to a lack of checkpoints?
    for i in range(3):
        nursery.start_soon(run_task, i)
    nursery.start_soon(saver, nursery, run_task, tmp_path / 'save.txt')
    assert nursery.cancel_scope.cancel_called


async def test_with_mock_clock_fixture(tmp_path, mock_clock):
    # Hangs indefinitely regardless of whether any `mock_clock.jump` is included
    async with trio.open_nursery() as nursery:
        mock_clock.jump(10)
        for i in range(3):
            nursery.start_soon(run_task, i)
        nursery.start_soon(saver, nursery, run_task, tmp_path / 'save.txt')
        mock_clock.jump(10)
    mock_clock.jump(10)
    assert nursery.cancel_scope.cancel_called


async def test_with_autojump_clock_fixture(tmp_path, autojump_clock):
    # PASSING: But how can I check that the nursery was cancelled at around 3 sec?
    async with trio.open_nursery() as nursery:
        for i in range(3):
            nursery.start_soon(run_task, i)
        nursery.start_soon(saver, nursery, run_task, tmp_path / 'save.txt')
    assert nursery.cancel_scope.cancel_called

Thanks in advance for your time and in answering my question! :slight_smile:

N.B. I understand that checking for a nursery cancellation circa 3 sec is unnecessarily strict/picky for this test, but I still would like to know if there’s a way to control the execution flow using time, so as to test more complex function logics that depend on time elapsed.

2

High-level thought: It sounds like you’re trying to run saver and the various run_task functions in the same nursery, which is resulting in an unnecessarily complex implementation. It’s also a less efficient one: the nursery will stay open up to 1 second after all the run_task tasks have completed, because saver only checks for that every second.

You might have better luck nesting multiple nurseries:

async def run_task(sleep_time: int) -> None:
    await trio.sleep(sleep_time)

async def saver(savepath: Path) -> None:
    """Save data periodically until cancelled."""
    while True:
        # This construct ensures you save one last time immediately after the
        # saver task is cancelled. If you don't care about that, you can just
        # save first and then sleep.
        try:
            await trio.sleep(1)
        finally:
            with open(savepath, "w") as f:
                f.write(str(trio.current_time()))

async def run_it() -> None:
    async with trio.open_nursery() as saver_nursery:
        saver_nursery.start_soon(saver, Path("example.txt"))
        async with trio.open_nursery() as tasks_nursery:
            for i in (1, 2, 3):
                tasks_nursery.start_soon(run_task, i)
        # Once the tasks_nursery block completes, all the tasks must be done,
        # so you can go ahead and cancel the saver
        saver_nursery.cancel_scope.cancel()

FWIW, this is a common pattern and there’s a concept in development of “service tasks” that will let you represent this in one nursery. That’s not committed to Trio yet though. See https://github.com/python-trio/trio/issues/1521 for details.

To answer your specific question about testing, this is a great place to use the autojump_clock fixture. It always starts at time 0 and is deterministic, so you can just assert trio.current_time() == 3.0 in your test_basic test. As a bonus, you won’t have to actually wait the three seconds when you run your tests.

In test_with_nursery_fixture, you’re seeing a failure because you never wait: you start the tasks, then immediately assert that the nursery was cancelled. But it hasn’t been cancelled yet, because the tasks haven’t run yet, and definitely haven’t yet run for three seconds.

In test_with_mock_clock_fixture, you’re using mock_clock rather than autojump_clock, so where you call jump() matters a lot. In this example, all of your calls to jump() are still before any of your tasks start running – you’ve started them, but no checkpoints have occurred, so they haven’t yet reached the sleep calls. Since you don’t have any jump() after reaching the sleep, you never complete your sleep, and thus the test appears to hang.

test_with_autojump_clock_fixture works correctly, and is indeed my suggestion here. You just need assert trio.current_time() == 3.0 at the end. :slight_smile:

1 Like
3

Thanks for your detailed response and the example usage of service nurseries — I didn’t think multi-nursery usage is a pattern at all!

Out of curiosity, would there be a way to make test_with_mock_clock_fixture work? I tried to add checkpoints by peppering that test with await trio.sleep(0) as per Trio’s docs but it didn’t work.

For completeness, here’s how I wrote my eventual test based on your suggestions and by using pytest-trio fixtures:

async def test_saver(tmp_path, nursery, autojump_clock):
    savepath = tmp_path / 'example.txt'
    nursery.start_soon(saver, savepath)  # service nursery
    async with trio.open_nursery() as tasks_nursery:
        for i in (1, 2, 3):
            tasks_nursery.start_soon(run_task, i)
    assert trio.current_time() == 3.
    assert savepath.exists()
4

Using MockClock without autojump is technically possible, but it’s quite cumbersome and rarely worthwhile. In this mode, time doesn’t advance unless you tell it to: adding checkpoints isn’t enough, you also have to somehow manage to call jump() at the right times. sleep() means “wait until N seconds from now”, so it won’t complete unless you jump() after the sleep begins. Keeping track of all of this for even a moderately involved test is too complex to seem tractable; I don’t recommend it.

5

I see! Thanks again for the explanations :slight_smile: