如何用PyTest测试FastAPI WebSocket连接是否已关闭？

问题背景 #

我正在用FastAPI WebSocket开发网络游戏后端通信机制，选用PyTest作为测试框架，目标是编写测试用例验证：当尝试连接不存在的房间时，连接会被立即关闭（管理器不允许连接不存在的房间）。但当前测试未通过，报错Failed: DID NOT RAISE <class 'starlette.websockets.WebSocketDisconnect'>。

现有代码问题分析 #

1. 同步测试的局限性：当前使用同步测试函数，但FastAPI的WebSocket操作是异步的，TestClient的同步websocket_connect会隐藏底层异步逻辑，无法正确捕获连接关闭的异常。
2. 异常被端点捕获：WebSocket端点中，manager.connect抛出的WebSocketDisconnect被try-except块捕获，客户端不会直接收到该异常，只会收到连接关闭的信号。
3. 上下文管理器的自动处理：client.websocket_connect的上下文管理器会自动处理连接关闭，不会将异常暴露给测试代码。

修复后的测试代码 #

改用异步测试，通过主动接收消息触发关闭检测：

import pytest
from fastapi.testclient import TestClient
from fastapi.websockets import WebSocketDisconnect

from src.game_manager import GameManager
from src.main import app, get_manager  # 确保导入get_manager依赖

async def override_manager() -> GameManager:
    try:
        yield override_manager.manager
    except AttributeError:
        manager = GameManager()
        manager.add_new_game(max_players=2, room_name="foo", password=None)
        manager.add_new_game(max_players=2, room_name="bar", password="123")
        override_manager.manager = manager
        yield override_manager.manager

app.dependency_overrides[get_manager] = override_manager

client = TestClient(app)

class TestWebsocketConnection:
    @pytest.mark.asyncio
    async def test_connect_to_non_existing_room(self):
        async with client.websocket_connect("/ws/non_existing_room") as ws:
            # 尝试接收消息，此时连接已被服务器关闭，会抛出WebSocketDisconnect
            with pytest.raises(WebSocketDisconnect):
                await ws.receive_json()

关键修改说明 #

• 异步测试适配：添加@pytest.mark.asyncio装饰器，使用异步上下文管理器async with，匹配FastAPI的异步运行特性。
• 主动触发异常检测：通过await ws.receive_json()尝试接收消息，此时服务器已关闭连接，会抛出WebSocketDisconnect，被pytest捕获完成断言。
• 依赖注入校验：确认导入get_manager，保证依赖覆盖逻辑正常生效。

替代验证方式 #

如果想更直观地断言结果，也可以用显式判断：

@pytest.mark.asyncio
async def test_connect_to_non_existing_room(self):
    async with client.websocket_connect("/ws/non_existing_room") as ws:
        try:
            await ws.receive_json()
        except WebSocketDisconnect:
            assert True
        else:
            assert False, "预期触发WebSocketDisconnect但未捕获"

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