[FIXED] How to make asyncio cancel() to actually cancel the task?

Issue

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?

Solution

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.

Answered By - VPfB