使用AsyncMock结合上下文管理器断言await调用失败求助（Python）

你的测试失败核心原因是对异步上下文管理器的mock逻辑理解有误：

• websockets.connect返回的是异步上下文管理器，客户端代码里的async with ... as websocket实际上会调用该管理器的__aenter__方法，最终拿到的websocket是这个方法的返回值，而非你直接创建的mock_websocket本身。
• 默认情况下，AsyncMock的__aenter__会生成一个新的AsyncMock实例，客户端中调用websocket.send的是这个新实例，而非你断言的mock_websocket，所以断言自然会提示调用次数为0。

解决方案一：让mock的__aenter__返回自身 #

修改测试代码，让异步上下文管理器返回你创建的mock对象，这样async with拿到的websocket就是你能追踪的实例：

import json
import unittest
from unittest.mock import AsyncMock, patch
from client import Client

class ConnectionTest(unittest.IsolatedAsyncioTestCase):
    async def test_message_is_sent(self):
        websocket_client = Client(interval=4)
        
        mock_websocket = AsyncMock()
        # 配置异步上下文管理器返回自身
        mock_websocket.__aenter__.return_value = mock_websocket

        with patch("client.websockets.connect", return_value=mock_websocket):
            await websocket_client.connect()
            message = {"messageId": "499"}
            mock_websocket.send.assert_awaited_once_with(json.dumps(message))

解决方案二：单独mock上下文管理器和websocket实例 #

拆分mock对象，明确指定上下文管理器的__aenter__返回websocket实例的mock：

import json
import unittest
from unittest.mock import AsyncMock, patch
from client import Client

class ConnectionTest(unittest.IsolatedAsyncioTestCase):
    async def test_message_is_sent(self):
        websocket_client = Client(interval=4)
        
        mock_connect = AsyncMock()
        mock_websocket = AsyncMock()
        # 指定connect返回的上下文管理器进入时返回websocket mock
        mock_connect.__aenter__.return_value = mock_websocket

        with patch("client.websockets.connect", return_value=mock_connect):
            await websocket_client.connect()
            message = {"messageId": "499"}
            mock_websocket.send.assert_awaited_once_with(json.dumps(message))

两种方案的本质都是确保async with拿到的websocket是你能追踪调用的mock实例，这样就能正确断言send方法的调用情况。

内容的提问来源于stack exchange，提问作者j1nrg