1. 理解FastAPI WebSocket与PyTest测试挑战
在使用FastAPI构建基于WebSocket的实时应用时，一个常见的需求是测试服务器在特定条件下主动关闭客户端连接的行为。例如，当客户端尝试连接到一个不存在的房间时，服务器应立即拒绝并关闭连接。PyTest是Python生态中流行的测试框架，结合FastAPI的TestClient，可以方便地对HTTP和WebSocket端点进行测试。
然而，测试WebSocket连接的关闭状态常常会遇到挑战。开发者可能会直观地尝试在建立连接的代码块外部使用pytest.raises(WebSocketDisconnect)来捕获异常，期望连接失败时立即抛出。然而，这种方法往往无法奏效，因为TestClient的websocket_connect方法可能成功建立底层TCP连接，但服务器端的WebSocket协议握手或业务逻辑处理随后导致连接关闭，此时异常并不会立即抛出。
考虑以下初始测试尝试及其返回的错误信息：
import pytest
from fastapi.testclient import TestClient
from fastapi.websockets import WebSocketDisconnect
# 假设app和get_manager以及override_manager已正确定义
# ... (省略了app和manager的依赖覆盖代码)
client = TestClient(app)
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 tests/test_websockets.py::TestWebsocketConnection::test_connect_to_non_existing_room - Failed: DID NOT RAISE 这个错误表明，尽管我们预期会抛出WebSocketDisconnect，但实际并没有。这是因为WebSocketDisconnect通常在尝试对一个已经关闭的WebSocket连接进行读写操作时才会触发，而不是在连接建立的瞬间。
2. WebSocketDisconnect异常的触发机制
WebSocketDisconnect是Starlette（FastAPI底层使用的Web框架）中定义的异常，它标志着WebSocket连接的意外断开或服务器主动关闭。理解其触发机制是编写有效测试的关键：
- 服务器主动关闭： 当服务器端代码调用websocket.close()方法，或者在处理连接过程中（例如在manager.connect方法中）抛出WebSocketDisconnect并被上层捕获后执行清理逻辑时，连接会被关闭。
- 客户端感知： 客户端通常不会在连接关闭的瞬间立即感知到异常。只有当客户端尝试通过已关闭的连接发送或接收数据时，底层网络库才会检测到连接状态的变化，并向上层抛出WebSocketDisconnect。
因此，要测试连接是否已关闭，我们需要模拟客户端尝试与服务器通信的场景。
3. 有效的测试策略：通过数据接收验证连接关闭
基于上述理解，测试WebSocket连接关闭的有效策略是：在尝试建立连接后，立即尝试从该WebSocket连接接收数据。如果服务器已经关闭了连接，那么这个接收数据的操作就会触发并抛出WebSocketDisconnect异常，我们就可以成功捕获它。
以下是实现这一策略的PyTest代码示例：
import pytest
from fastapi.testclient import TestClient
from fastapi.websockets import WebSocketDisconnect
from typing import Annotated
# 假设你的FastAPI应用和GameManager的定义如下
# src/game_manager.py
class GameManager:
def __init__(self):
self.games = {} # 存储游戏房间信息
async def connect(self, websocket, room_name, password):
if room_name not in self.games:
# 如果房间不存在，则抛出WebSocketDisconnect
raise WebSocketDisconnect(code=1008, reason="Room does not exist")
# 实际连接逻辑...
await websocket.accept()
print(f"Client connected to room: {room_name}")
# 这里为了测试，假设连接成功后不会立即发送数据
async def remove(self, websocket):
# 清理连接逻辑
print("Client disconnected.")
async def handle_message(self, room_name, client_id, data):
# 处理消息逻辑
pass
# src/main.py
from fastapi import FastAPI, APIRouter, Depends, WebSocket
from fastapi.routing import APIRoute
# 为了演示，这里简化get_manager
def get_manager() -> GameManager:
return GameManager()
app = FastAPI()
router = APIRouter()
@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)
# client_id = websocket.scope["client_id"] # 实际应用中会获取
while True:
data = await websocket.receive_json()
# await manager.handle_message(room_name, client_id, data) # 实际应用中会处理
except WebSocketDisconnect:
await manager.remove(websocket)
except Exception as e:
print(f"Unexpected error: {e}")
await manager.remove(websocket)
app.include_router(router)
# tests/test_websockets.py
# 依赖覆盖，确保测试环境隔离且可控
async def override_get_manager() -> GameManager:
try:
# 尝试使用已存在的manager实例
yield override_get_manager.manager
except AttributeError:
# 如果不存在，则创建并初始化一个新的manager
manager = GameManager()
manager.games["foo"] = {} # 添加一个存在的房间用于其他测试
override_get_manager.manager = manager
yield override_get_manager.manager
# 将依赖覆盖应用到FastAPI应用
app.dependency_overrides[get_manager] = override_get_manager
client = TestClient(app)
class TestWebsocketConnection:
def test_connect_to_non_existing_room_correctly_closed(self):
"""
测试连接到不存在的房间时，连接是否被正确关闭。
通过尝试接收数据来触发WebSocketDisconnect异常。
"""
with pytest.raises(WebSocketDisconnect) as excinfo:
with client.websocket_connect("/ws/non_existing_room") as ws:
# 关键步骤：尝试从已关闭的连接接收数据
# 这将触发并捕获WebSocketDisconnect异常
ws.receive_json()
# 可选：进一步断言异常的详细信息，例如错误码或原因
assert excinfo.type is WebSocketDisconnect
assert excinfo.value.code == 1008
assert "Room does not exist" in excinfo.value.reason在这个示例中，ws.receive_json()是关键。当客户端尝试连接到/ws/non_existing_room时，服务器端的manager.connect方法会检测到房间不存在，并立即抛出WebSocketDisconnect。websocket_endpoint捕获此异常后，会执行清理逻辑（manager.remove），但不会向客户端发送任何数据。此时，客户端的WebSocket连接实际上已经被服务器关闭。当客户端代码执行到ws.receive_json()时，由于连接已关闭，它会检测到这一点并抛出WebSocketDisconnect，从而被pytest.raises成功捕获。
4. 注意事项与最佳实践
- 服务器端行为： 确保服务器端在需要关闭连接时，要么显式调用websocket.close()，要么通过抛出WebSocketDisconnect并被上层捕获来间接导致连接关闭。客户端的测试方法依赖于服务器的这种行为。
- 依赖注入覆盖： 在测试中，使用app.dependency_overrides来替换真实的GameManager实例，可以确保测试环境的隔离性和可控性。这允许你在每个测试中为GameManager设置不同的初始状态，例如预设存在的房间。
- 异常细节断言： 除了捕获WebSocketDisconnect本身，还可以进一步断言异常对象的code和reason属性，以验证连接关闭的原因是否符合预期。这增加了测试的健壮性。
- 避免死循环： 在服务器端的websocket_endpoint中，如果manager.connect成功，通常会进入一个while True循环来持续接收消息。但在测试连接关闭的场景中，如果manager.connect失败并抛出异常，这个循环就不会被执行，这正是我们期望的行为。
5. 总结
通过理解WebSocketDisconnect异常的触发时机，并采用在连接建立后尝试接收数据的策略，我们可以有效地在FastAPI应用中使用PyTest测试WebSocket连接的关闭情况。这种方法不仅能够准确捕获预期的异常，还能帮助开发者验证服务器端在特定业务逻辑下对WebSocket连接的正确管理。
