如何测试含嵌套async with的异步代码？Mock await报错求助

你遇到的TypeError: object MagicMock can't be used in 'await' expression错误，本质是普通MagicMock没有实现异步操作所需的协议：异步上下文管理器需要__aenter__/__aexit__异步方法，可等待对象需要__await__方法，而默认的MagicMock不具备这些特性，无法被await或用于async with语句。

推荐使用Python 3.8+自带的AsyncMock（专门用于模拟异步对象），它原生支持异步上下文管理器和异步方法，代码更简洁易维护。以下是完整的测试用例实现：

from unittest import TestCase
from unittest.mock import patch, AsyncMock
import asyncio

class TestClass(TestCase):
    def test_function(self):
        # 定义预设的返回数据
        test_rows = [('row1_data',), ('row2_data',)]
        # 如果需要验证SQL语句，可提前定义预期值
        expected_sql = "SELECT * FROM your_table WHERE ..."

        with patch('some_direction.session_pool') as mock_session_pool:
            # 1. 模拟session_pool返回的连接池对象
            mock_pool = AsyncMock()
            mock_session_pool.return_value = mock_pool

            # 2. 模拟pool.acquire()异步上下文管理器，返回连接对象conn
            mock_conn = AsyncMock()
            mock_pool.acquire.return_value.__aenter__.return_value = mock_conn

            # 3. 模拟conn.cursor()异步上下文管理器，返回游标对象cursor
            mock_cursor = AsyncMock()
            mock_conn.cursor.return_value.__aenter__.return_value = mock_cursor

            # 4. 模拟cursor.execute()异步方法，返回带fetchall的结果对象
            mock_res = AsyncMock()
            mock_res.fetchall.return_value = test_rows
            mock_cursor.execute.return_value = mock_res

            # 执行测试逻辑
            result = asyncio.run(ClosingFunds.find(some_values))

            # 验证结果和方法调用
            self.assertEqual(result, test_rows)
            mock_session_pool.assert_called_once()
            mock_pool.acquire.assert_called_once()
            mock_conn.cursor.assert_called_once()
            mock_cursor.execute.assert_called_once_with(expected_sql)  # 可选：验证SQL参数

低Python版本兼容方案（<3.8） #

如果你的Python版本低于3.8，没有AsyncMock，可以手动给MagicMock添加异步协议支持：

from unittest import TestCase
from unittest.mock import patch, MagicMock
import asyncio

class TestClass(TestCase):
    def test_function(self):
        test_rows = [('row1_data',), ('row2_data',)]

        with patch('some_direction.session_pool') as mock_session_pool:
            # 模拟session_pool为异步函数
            mock_pool = MagicMock()
            mock_session_pool.return_value.__await__ = lambda self: iter([mock_pool])

            # 模拟acquire异步上下文管理器
            mock_conn = MagicMock()
            mock_pool.acquire.return_value.__aenter__ = MagicMock(
                return_value=asyncio.coroutine(lambda: mock_conn)()
            )

            # 模拟cursor异步上下文管理器
            mock_cursor = MagicMock()
            mock_conn.cursor.return_value.__aenter__ = MagicMock(
                return_value=asyncio.coroutine(lambda: mock_cursor)()
            )

            # 模拟execute异步方法
            mock_res = MagicMock()
            mock_res.fetchall.return_value = test_rows
            mock_cursor.execute = MagicMock(
                return_value=asyncio.coroutine(lambda: mock_res)()
            )

            result = asyncio.run(ClosingFunds.find(some_values))
            self.assertEqual(result, test_rows)

内容的提问来源于stack exchange，提问作者Matias Ezequiel Miguel Rovedo