In Python, how do I test a function that makes asyncio.ensure_future(...) calls?

user_2023861 0
Viewed 1047

In python unit tests, how do I assert that a function was called if that function is called using asyncio.ensure_future(...)? I have this simple test that fails:

async def test_assert_ensure_future_called():
    async def increment():
        increment.call_count += 1

    increment.call_count = 0

    asyncio.ensure_future(increment())
    asyncio.ensure_future(increment())
    asyncio.ensure_future(increment())

    # await asyncio.sleep(0)  # this allows the asyncio.ensure_future calls to run
    assert increment.call_count == 3  # this fails because the calls to increment() haven't happened yet

I found that if I insert a call like await asyncio.sleep(0) just before the assert statement, that the test will succeed. I assume that this works because it deprioritizes the test task and lets the calls to increment go first. Is this the proper way to do this in a test though?

future python-asyncio python-3.x unit-testing
1 Answers

Save references of those tasks and await it either calling await per tasks or using asyncio.gather for one-line.

import asyncio


async def test_assert_ensure_future_called():
    async def increment():
        increment.call_count += 1

    increment.call_count = 0


    tasks = [asyncio.ensure_future(increment()) for _ in range(3)]
    await asyncio.gather(*tasks)

    assert increment.call_count == 3


asyncio.run(test_assert_ensure_future_called())
user_109090290
Related