背景与挑战
在开发基于fastapi的websocket应用程序时，尤其是在构建需要严格连接管理的游戏或其他实时通信服务时，测试服务器端主动关闭连接的场景至关重要。例如，当客户端尝试连接到一个不存在的房间时，服务器理应立即拒绝并关闭该连接。
然而，在使用PyTest和fastapi.testclient.TestClient进行测试时，直接在client.websocket_connect上下文管理器外部或内部尝试捕获WebSocketDisconnect异常，往往会发现测试失败，并提示“DID NOT RAISE
FastAPI WebSocket关闭机制解析
理解WebSocketDisconnect异常的触发机制是解决问题的关键。在FastAPI（底层基于Starlette）中，WebSocketDisconnect通常表示WebSocket连接已经断开或被关闭。这个异常不会在连接建立（即HTTP握手成功）的那一刻立即抛出，即使服务器端立即决定关闭连接。相反，它会在客户端尝试对一个已经关闭的WebSocket连接执行读（receive_json(), receive_text(), receive_bytes()）或写（send_json(), send_text(), send_bytes()）操作时被检测到并抛出。
当服务器端通过websocket.close()方法显式关闭连接，或者在处理逻辑中（例如，通过依赖注入的管理器）判断连接不合法并主动抛出WebSocketDisconnect时，连接状态会变为关闭。此时，客户端如果尝试继续与该连接交互，便会收到WebSocketDisconnect异常。
PyTest测试实践：验证连接关闭
为了有效测试服务器端主动关闭WebSocket连接的行为，我们需要模拟客户端在连接关闭后尝试与服务器通信的场景。
错误尝试及原因分析
最初的尝试可能是在client.websocket_connect上下文管理器外部或内部直接使用pytest.raises(WebSocketDisconnect)：
# 错误的测试方法示例
import pytest
from fastapi.testclient import TestClient
from fastapi.websockets import WebSocketDisconnect
# ... 其他设置和依赖覆盖 ...
class TestWebsocketConnection:
def test_connect_to_non_existing_room_incorrect(self):
with pytest.raises(WebSocketDisconnect) as e_info:
with client.websocket_connect("/ws/non_existing_room") as ws:
# 即使服务器立即关闭，这里也可能不立即抛出
ws.send_json({"message": "Hello world"}) # 尝试发送数据这种方法失败的原因在于，client.websocket_connect主要负责完成WebSocket的HTTP握手过程。即使服务器端在握手完成后立即通过业务逻辑（如manager.connect）决定关闭连接，这种关闭行为通常不会立即导致websocket_connect本身抛出WebSocketDisconnect。客户端只有在尝试发送或接收数据时，才会发现底层连接已经断开，从而触发WebSocketDisconnect。
正确的测试方法
正确的策略是在连接建立后，立即尝试从WebSocket接收数据。如果服务器已经关闭了连接，那么这个接收操作将不可避免地失败并抛出WebSocketDisconnect。
import pytest
from fastapi.testclient import TestClient
from fastapi.websockets import WebSocketDisconnect
from typing import Annotated
# 假设你的FastAPI应用和GameManager都在src目录下
from src.game_manager import GameManager
from src.main import app, get_manager # 假设get_manager是你的依赖函数
# 依赖覆盖：为测试环境提供一个预设的GameManager实例
# 这确保了测试环境中的GameManager状态可控
async def override_manager() -> GameManager:
try:
# 如果manager已经存在，直接使用
yield override_manager.manager
except AttributeError:
# 否则，创建一个新的GameManager实例并添加测试房间
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
# 将测试用的依赖覆盖到FastAPI应用中
app.dependency_overrides[get_manager] = override_manager
# 初始化FastAPI测试客户端
client = TestClient(app)
class TestWebsocketConnection:
def test_connect_to_non_existing_room_and_expect_disconnect(self):
"""
测试连接到不存在的房间时，服务器应立即关闭连接。
客户端尝试接收数据时，应检测到连接已关闭并抛出WebSocketDisconnect。
"""
with pytest.raises(WebSocketDisconnect):
with client.websocket_connect("/ws/non_existing_room") as ws:
# 服务器在连接到不存在的房间时，会通过其内部逻辑（如manager.connect）
# 拒绝连接或立即关闭。
# 客户端在此处尝试接收数据时，会检测到连接已关闭，
# 从而抛出WebSocketDisconnect。
ws.receive_json()
在上述代码中，with client.websocket_connect("/ws/non_existing_room") as ws:语句会尝试建立连接。即使服务器端的manager.connect方法立即抛出WebSocketDisconnect（如问题描述所述），导致服务器端关闭连接，客户端的websocket_connect在这一步可能不会立即感知到。但当客户端执行ws.receive_json()时，它会尝试从一个已经关闭的套接字读取数据，此时WebSocketDisconnect异常就会被抛出，并被pytest.raises成功捕获，从而验证了连接关闭的预期行为。
注意事项
- 异常触发时机： WebSocketDisconnect异常是在客户端尝试与一个已关闭的WebSocket连接进行I/O操作（发送或接收数据）时触发的，而不是在连接建立或服务器端关闭连接的瞬间。
- 服务器端逻辑： 确保你的FastAPI WebSocket处理函数（例如websocket_endpoint中的manager.connect部分）确实在特定条件下执行了关闭连接的操作，无论是通过websocket.close()还是通过抛出WebSocketDisconnect。
- 依赖覆盖： 在测试中，使用app.dependency_overrides来替换生产环境的依赖项（如GameManager），可以让你在测试中完全控制依赖的行为，模拟各种场景。
- 客户端行为： 此方法适用于测试服务器主动关闭连接的场景。如果客户端主动关闭连接，或者连接因网络问题而断开，WebSocketDisconnect也会被触发。
总结
通过在PyTest中利用pytest.raises(WebSocketDisconnect)并结合在client.websocket_connect后立即尝试调用ws.receive_json()（或其他接收数据的方法），我们可以有效地测试FastAPI WebSocket连接在服务器端被主动关闭的场景。这种方法精准地捕获了连接断开的信号，为WebSocket应用程序的健壮性提供了可靠的测试保障。
