1. 理解WebSocket连接关闭与测试需求
在开发基于fastapi的websocket应用时，我们经常需要处理各种连接状态，包括正常的建立、通信以及异常的断开。特别是在某些业务逻辑下，服务器可能会主动拒绝或关闭客户端的连接，例如当客户端尝试连接一个不存在的房间或未授权的资源时。为了确保这些断开逻辑的正确性，编写可靠的测试用例至关重要。
使用Pytest作为测试框架，FastAPI提供了TestClient来模拟客户端行为，包括WebSocket连接。我们的目标是验证当服务器主动关闭连接时，客户端能够正确地捕获到WebSocketDisconnect异常。
2. 初始尝试及遇到的问题
最初的测试思路可能是直接将pytest.raises(WebSocketDisconnect)上下文管理器包裹在client.websocket_connect()调用外部，期望连接失败时立即捕获异常。
import pytest
from fastapi.testclient import TestClient
from fastapi.websockets import WebSocketDisconnect
# 假设 app 和 get_manager, override_manager 等已定义
# ... (省略依赖注入和TestClient初始化代码) ...
class TestWebsocketConnection:
def test_connect_to_non_existing_room_initial_attempt(self):
with pytest.raises(WebSocketDisconnect) as e_info:
with client.websocket_connect("/ws/non_existing_room") as ws:
# 尝试发送数据，但可能此时连接尚未完全关闭或异常未传播
ws.send_json({"message": "Hello world"})然而，这种方法通常会导致测试失败，并报错FAILED: DID NOT RAISE <class 'starlette.websockets.WebSocketDisconnect'>。这是因为client.websocket_connect()在某些情况下可能会成功建立底层的TCP连接，即使服务器端的业务逻辑立即决定关闭WebSocket连接。WebSocketDisconnect异常通常不是在连接建立的瞬间抛出，而是在客户端尝试对一个已经关闭的WebSocket连接执行操作（如接收或发送数据）时才会显现。
在服务器端，当manager.connect()方法在检测到无效房间时抛出WebSocketDisconnect，FastAPI的WebSocket路由会捕获这个异常，并相应地关闭连接。
@router.websocket("/ws/{room_name}")
@router.websocket("/ws/{room_name}/{password}")
async def websocket_endpoint(
websocket: WebSocket,
manager: Annotated[GameManager, Depends(get_manager)],
):
room_name = websocket.path_params["room_name"]
password = websocket.path_params.get("password", None)
try:
await manager.connect(websocket, room_name, password)
# 如果 manager.connect 成功，客户端才会被接受并进入消息循环
client_id = websocket.scope["client_id"]
while True:
data = await websocket.receive_json()
await manager.handle_message(room_name, client_id, data)
except WebSocketDisconnect:
# 当 manager.connect 抛出 WebSocketDisconnect 或客户端断开时，这里会捕获
await manager.remove(websocket)
# FastAPI/Starlette 框架会确保连接被正确关闭3. 正确测试WebSocket连接关闭的方法
解决上述问题的关键在于理解WebSocketDisconnect的触发时机。它通常在客户端尝试从一个已关闭或断开的WebSocket连接接收数据时被抛出。因此，在客户端成功连接（即使服务器立即关闭）后，我们应该主动尝试从连接中接收数据。
3.1 核心思路
- 客户端使用client.websocket_connect()尝试连接。
- 即使服务器端因业务逻辑立即关闭连接，websocket_connect本身可能不会立即抛出WebSocketDisconnect。
- 在with client.websocket_connect(...) as ws:块内部，立即调用ws.receive_json()（或ws.receive_text()、ws.receive_bytes()）。
- 如果服务器已经关闭了连接，尝试接收数据将触发WebSocketDisconnect异常。
- 将整个过程包裹在pytest.raises(WebSocketDisconnect)中，即可成功捕获并验证异常。
3.2 示例代码
以下是经过修正的、能够正确测试WebSocket连接关闭的Pytest用例：
import pytest
from fastapi.testclient import TestClient
from fastapi.websockets import WebSocketDisconnect
from typing import Annotated
# 假设以下导入和定义已存在
from src.game_manager import GameManager
from src.main import app
from fastapi import Depends, WebSocket, APIRouter # 假设 router 是一个 APIRouter 实例
# 假设 get_manager 是一个依赖函数
def get_manager():
# 这是一个占位符，实际应用中会从依赖注入系统获取 GameManager 实例
pass
# 重写依赖以进行测试
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:
def test_connect_to_non_existing_room(self):
"""
测试连接到不存在的房间时，WebSocket连接是否被正确关闭。
"""
with pytest.raises(WebSocketDisconnect):
with client.websocket_connect("/ws/non_existing_room") as ws:
# 尝试从已关闭的连接接收数据，这将触发 WebSocketDisconnect
ws.receive_json()
def test_connect_to_existing_room(self):
"""
测试连接到存在的房间时，连接是否成功建立且未断开。
"""
with client.websocket_connect("/ws/foo") as ws:
# 确保连接成功，可以发送和接收数据
ws.send_json({"message": "Hello from foo"})
# 接收一条消息，验证连接活跃
response = ws.receive_json()
assert "type" in response and response["type"] == "websocket.send" # 示例断言
# 确保连接在退出 with 块后正常关闭
with pytest.raises(WebSocketDisconnect):
# 尝试在连接关闭后接收数据，验证其已关闭
ws.receive_json()3.3 注意事项
- 服务器端行为是关键： 这种测试方法的前提是服务器端确实在某些条件下（如manager.connect失败）主动关闭了WebSocket连接。如果服务器只是不发送数据但保持连接开放，receive_json()将可能无限期等待或超时，而不是抛出WebSocketDisconnect。
- WebSocketDisconnect的来源： WebSocketDisconnect异常可能由多种情况触发，包括客户端主动断开、服务器主动断开、网络问题等。在测试中，我们确保它是由于服务器端业务逻辑拒绝连接导致的。
- 依赖注入： 在测试中，通常需要通过app.dependency_overrides来模拟或控制依赖项（如GameManager）的行为，以确保测试的隔离性和可预测性。
- 测试覆盖： 除了测试连接关闭，还应测试正常连接、消息发送与接收、以及其他异常情况。
4. 总结
在Pytest中测试FastAPI WebSocket连接的关闭，尤其是当服务器因业务逻辑主动断开连接时，关键在于模拟客户端尝试从一个已关闭的连接接收数据。通过将ws.receive_json()（或类似接收方法）调用包裹在pytest.raises(WebSocketDisconnect)上下文管理器中，我们可以有效地验证服务器是否按预期关闭了连接。这种方法能够精确地捕获到WebSocketDisconnect异常，从而确保WebSocket连接管理逻辑的健壮性。
