How to make asyncio cancel() to actually cancel the task?

Question

This code is stuck in an infinite loop, the self.task.cancel() seems to have no effect:

import asyncio

from unittest.mock import AsyncMock, patch


async def loop():
    while True:
        await asyncio.sleep(1)


class AsyncSleepMock(AsyncMock):
    def __init__(self):
        super(AsyncMock, self).__init__()
        self.task = None

    async def __call__(self, delay, *args, **kwargs):
        self.task.cancel()
        return await super(AsyncMock, self).__call__(delay, *args, **kwargs)


def create_async_sleep_mock():
    return AsyncSleepMock()


@patch("asyncio.sleep", new_callable=create_async_sleep_mock)
def main(async_sleep_mock):
    loop_task = asyncio.get_event_loop().create_task(loop())
    async_sleep_mock.task = loop_task
    asyncio.get_event_loop().run_until_complete(loop_task)


if __name__ == "__main__":
    main()

The goal is to make a mock of asyncio.sleep() that can break out from that infinite loop() that the application under test has. How to do that?

Answer 1

self.task.cancel() marks the task as cancelled, but at that moment is this the active task on CPU. A task switch must occur to allow the scheduler to process the cancellation.

From the cancel() docs:

Request the Task to be cancelled.

This arranges for a CancelledError exception to be thrown into the wrapped coroutine on the next cycle of the event loop.

I have inserted an unmocked await asyncio.sleep(0) to ensure the needed task switch, now it doesn't loop any more:

realsleep = asyncio.sleep

class AsyncSleepMock(AsyncMock):
    def __init__(self):
        super(AsyncMock, self).__init__()
        self.task = None

    async def __call__(self, delay, *args, **kwargs):
        self.task.cancel()
        await realsleep(0)
        return await super(AsyncMock, self).__call__(delay, *args, **kwargs)

---

For completness, I'm adding a quote from the asyncio.sleep() description:

sleep() always suspends the current task, allowing other tasks to run.

Setting the delay to 0 provides an optimized path to allow other tasks to run. This can be used by long-running functions to avoid blocking the event loop for the full duration of the function call.