asyncio and pytest: get all tasks from event loop of tested function

116 views Asked by At

I have fetcher.py module with code below and I need to test it. I use pytest with pytest-asyncio plugin for that. How to get all tasks from event loop of tested function in tests.py module? When I try this code:
tests.py

import asyncio
import pytest
from fetcher import main

    @pytest.mark.asyncio
    async def test_01(self):
        await main(10)
        all_tasks = len(asyncio.all_tasks()) # get number of tasks
        assert all_tasks == 10

I got an assertion error:

1 != 10

Expected :10
Actual   :1

I think its because tasks are taken from pytest event loop, not from main(10) event loop. And there is only 1 task.

fetcher.py

import asyncio


async def worker(queue: asyncio.Queue):
    while True:
        current_url = await queue.get()
        try:
            if current_url is None:
                break
            print(f"{current_url}: OK")
            print(len(asyncio.all_tasks())) # prints 10
        finally:
            queue.task_done()


async def main(max_requests):
    queue = asyncio.Queue(maxsize=max_requests)
    workers = {asyncio.create_task(worker(queue)) for _ in range(max_requests)}
    urls = ["https://example1.org/", "https://example2.org/", "https://example3.org/", "https://example4.org/"]
    for url in urls:
        await queue.put(url)
    for _ in range(len(workers)):
        await queue.put(None)
    await queue.join()

if __name__ == "__main__":
    asyncio.run(main(10))
1

There are 1 answers

2
Andrew Svetlov On

I think the test result is correct.

Initially, the code has 1 task: the main task that executes test_01 code.

After that, the test calls await main(10). Inside main() async function creates 10 additional tasks. The total task number is 11 (main + 10 workers).

Before main() exists it asks worker tasks to stop by sending None into the queue and waits for its processing. Now the number of tasks is 1 again.

At the moment of assert all_tasks == 10 execution the number of tasks is 1. Exactly as the test says.