如何在 Python 測試中模擬 Awaitable

Question

我想知道是否有一種方法可以輕松地模擬用於 python 測試的Awaitable object。 我知道如何使用AsyncMock來創建鏡像協程的模擬（即將從方法調用返回Awaitable對象），但是在嘗試直接模擬Awaitable時遇到了一些問題。

這是一個例子：

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()

拋出：

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

深入研究 Python 文檔，看起來__await__魔術方法是模擬中不支持的少數幾個方法之一，所以想知道我怎么能 go 這樣做？

Answer 1

一種解決方案可能是將 AsyncMock 擴展為可等待的模擬 object，添加await 。 在代碼中，這可能看起來像

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