1. 引言：FastAPI WebSocket连接关闭测试的重要性
在开发基于FastAPI的WebSocket应用时，确保连接管理逻辑的健壮性至关重要。特别是在某些业务场景下，服务器可能需要根据客户端的请求条件（例如，尝试连接到不存在的资源或未经授权的访问）立即拒绝并关闭WebSocket连接。为了验证这些服务器端的主动关闭行为是否按预期工作，我们需要编写有效的单元测试。
FastAPI配合PyTest是常见的测试组合，但测试WebSocket连接的关闭状态有时会遇到挑战。传统的异常捕获方式可能无法直接捕捉到WebSocket连接的断开事件，尤其是在连接建立过程中或建立后立即断开的情况下。
2. 问题分析：初始尝试与困境
许多开发者在尝试测试WebSocket连接关闭时，可能会直观地将pytest.raises(WebSocketDisconnect)用于包裹整个WebSocket连接建立的代码块，例如：
import pytest
from fastapi.testclient import TestClient
from fastapi.websockets import WebSocketDisconnect
# ... (FastAPI应用和依赖覆盖设置) ...
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: DID NOT RAISE
在上述场景中，如果服务器端在manager.connect方法中判断房间不存在并立即抛出WebSocketDisconnect（或者调用websocket.close()而不进行websocket.accept()），client.websocket_connect可能会成功返回一个WebSocket对象，但该对象背后代表的连接实际上已经被服务器关闭。此时，直接在with client.websocket_connect(...)块外部捕获异常是无效的，因为异常尚未在客户端代码中向上冒泡。
3. 解决方案：利用 receive_json() 触发断开异常
解决这个问题的关键在于理解WebSocketDisconnect异常的触发机制。当一个WebSocket连接在服务器端被关闭后，客户端尝试从该连接接收数据时，会感知到连接已关闭的状态，并抛出WebSocketDisconnect异常。我们可以利用这一特性来验证连接的关闭。
核心思路：
- 使用client.websocket_connect尝试建立连接。
- 在连接建立后（即使服务器已立即关闭），尝试从该WebSocket对象接收数据（例如，ws.receive_json()）。
- 如果服务器确实关闭了连接，此接收操作将立即触发WebSocketDisconnect异常。
- 使用pytest.raises(WebSocketDisconnect)来捕获并断言这个预期异常。
以下是实现此方案的测试代码：
import pytest
from fastapi.testclient import TestClient
from fastapi.websockets import WebSocketDisconnect
from typing import Annotated
from fastapi import Depends, APIRouter, WebSocket
# 假设你的应用结构如下，这里为了完整性提供示例
# src/game_manager.py
class GameManager:
def __init__(self):
self.games = {}
def add_new_game(self, max_players, room_name, password):
self.games[room_name] = {"max_players": max_players, "password": password, "clients": {}}
async def connect(self, websocket: WebSocket, room_name: str, password: str | None):
if room_name not in self.games:
# 关键：如果房间不存在，立即抛出 WebSocketDisconnect
raise WebSocketDisconnect(code=1008, reason="Room does not exist")
# 假设这里会处理密码验证等，并最终接受连接
await websocket.accept()
# 假设 client_id 是从某个地方生成的
client_id = f"client_{len(self.games[room_name]['clients'])}"
websocket.scope["client_id"] = client_id
self.games[room_name]["clients"][client_id] = websocket
print(f"Client {client_id} connected to {room_name}")
async def handle_message(self, room_name, client_id, data):
print(f"Received message from {client_id} in {room_name}: {data}")
async def remove(self, websocket: WebSocket):
# 实际的移除逻辑
print(f"Client {websocket.scope.get('client_id')} disconnected.")
# src/main.py
from fastapi import FastAPI
app = FastAPI()
router = APIRouter()
def get_manager():
# 实际应用中可能是单例或依赖注入
return GameManager()
@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)
app.include_router(router)
# tests/test_websockets.py
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
# 假设 get_manager 是你的依赖注入函数
from src.main import get_manager
app.dependency_overrides[get_manager] = override_manager
client = TestClient(app)
class TestWebsocketConnection:
def test_connect_to_non_existing_room_solution(self):
# 使用 pytest.raises 包裹，并在连接建立后尝试接收数据
with pytest.raises(WebSocketDisconnect):
with client.websocket_connect("/ws/non_existing_room") as ws:
# 关键步骤：尝试从已关闭的连接接收数据
ws.receive_json()在这个修正后的测试中，当client.websocket_connect("/ws/non_existing_room")被调用时，服务器端的manager.connect方法会因为房间不存在而抛出WebSocketDisconnect。由于这个异常发生在websocket.accept()之前，websocket.accept()不会被调用，连接实际上未完全建立。客户端的TestClient在此时可能仍然会返回一个WebSocket对象。然而，当客户端随后尝试调用ws.receive_json()时，它会发现底层连接已经关闭，从而在客户端侧抛出WebSocketDisconnect异常。pytest.raises就能成功捕获到这个异常，验证了服务器正确地拒绝了连接。
4. 服务器端行为的配合
为了使上述测试方法有效，服务器端的WebSocket处理逻辑必须与预期行为相符：
- 如果服务器不应允许连接： 服务器应在调用websocket.accept()之前，通过抛出WebSocketDisconnect异常或直接调用websocket.close()来拒绝连接。在我们的示例中，manager.connect方法在检测到房间不存在时直接抛出了WebSocketDisconnect，这正是测试所期望的。
- 如果服务器在接受后立即关闭： 即使服务器调用了websocket.accept()，但随后立即调用websocket.close()且不发送任何数据，客户端的ws.receive_json()仍然会捕获到WebSocketDisconnect。
重要的是，服务器在关闭连接时，不应在关闭前发送任何JSON消息，否则ws.receive_json()可能会先接收到一条消息，而不是直接抛出异常。
5. 总结与最佳实践
测试FastAPI WebSocket连接的关闭状态，特别是当服务器主动拒绝连接时，需要巧妙地利用WebSocketDisconnect异常的触发时机。
- 核心策略： 在预期连接会立即关闭的场景下，建立连接后，尝试从WebSocket客户端接收数据（如ws.receive_json()）。如果连接确实被服务器关闭，此操作将导致客户端抛出WebSocketDisconnect异常。
- 服务器端配合： 确保服务器在拒绝连接时，要么在websocket.accept()之前抛出WebSocketDisconnect，要么在websocket.accept()之后立即关闭连接且不发送任何消息。
- PyTest用法： 使用with pytest.raises(WebSocketDisconnect):来精确捕获并验证这个预期的异常。
通过这种方法，我们可以编写出健壮且准确的测试，以验证FastAPI WebSocket应用在各种异常情况下的连接管理逻辑。
