[FIXED] How to mock an Awaitable in Python tests

Issue

I'm wondering if there's a way to easily mock an Awaitable object for python tests. I know how to use AsyncMocks to create mocks which mirror coroutines (i.e. will return Awaitable objects from method calls), but am running into a bit of problem when trying to mock an Awaitable directly.

Here's an example:

import unittest
from asyncio import Task
from typing import Any
from unittest import IsolatedAsyncioTestCase
from unittest.mock import AsyncMock


async def function_using_awaitable(task: Task) -> Any:
    await task
    return task.result()

class TestFunction(IsolatedAsyncioTestCase):
    async def test_function_using_awaitable(self):
        mock_task = AsyncMock(spec=Task)

        result = await function_using_awaitable(mock_task)

        mock_task.assert_awaited_once()
        assert result == mock_task.result.return_value

if __name__ == '__main__':
    unittest.main()

which throws up:

Traceback (most recent call last):
  ...
  File ..., line 9, in function_using_awaitable
    await task
TypeError: object AsyncMock can't be used in 'await' expression

Digging into the Python docs it looks like the __await__ magic method is one of the few that isn't supported in mocks so wondering how I can go about doing this?

Solution

One solution might be to just extend the AsyncMock to be an awaitable mock object, adding await. In code, that could look like

from typing import Iterator, Any

async def function_using_awaitable(task: Task) -> Any:
    await task
    return task.result()

class AwaitableMock(AsyncMock):
    def __await__(self) -> Iterator[Any]:
        self.await_count += 1
        return iter([])

class TestFunction(IsolatedAsyncioTestCase):
    async def test_function_using_awaitable(self):
        mock_task = AwaitableMock(spec=Task)

        result = await function_using_awaitable(mock_task)

        mock_task.assert_awaited_once()
        assert result == mock_task.result.return_value

Answered By - Simon Hawe