本文探讨了在FastAPI应用中使用PyTest测试WebSocket连接异常关闭的方法。当服务器端基于业务逻辑（如连接到不存在的房间）主动关闭连接时，客户端测试需要验证此行为。通过在连接建立后尝试从已关闭的WebSocket接收数据，可以有效捕获并断言WebSocketDisconnect异常，从而确保连接关闭逻辑的正确性。
1. 引言：FastAPI WebSocket连接关闭测试的重要性
在开发基于FastAPI的WebSocket应用时，确保连接管理逻辑的健壮性至关重要。特别是在某些业务场景下，服务器可能需要根据客户端的请求条件（例如，尝试连接到不存在的资源或未经授权的访问）立即拒绝并关闭WebSocket连接。为了验证这些服务器端的主动关闭行为是否按预期工作，我们需要编写有效的单元测试。
FastAPI配合PyTest是常见的测试组合，但测试WebSocket连接的关闭状态有时会遇到挑战。传统的异常捕获方式可能无法直接捕捉到WebSocket连接的断开事件，尤其是在连接建立过程中或建立后立即断开的情况下。
2. 问题分析：初始尝试与困境
许多开发者在尝试测试WebSocket连接关闭时，可能会直观地将pytest.raises(WebSocketDisconnect)用于包裹整个WebSocket连接建立的代码块，例如：
import pytestfrom fastapi.testclient import TestClientfrom 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 websockets.WebSocketDisconnect’>。其原因在于，WebSocketDisconnect异常通常不是在client.websocket_connect调用时立即抛出，而是在客户端尝试与一个已经断开的连接进行通信（如发送或接收数据）时才会被感知并抛出。
在上述场景中，如果服务器端在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 pytestfrom fastapi.testclient import TestClientfrom fastapi.websockets import WebSocketDisconnectfrom typing import Annotatedfrom fastapi import Depends, APIRouter, WebSocket# 假设你的应用结构如下，这里为了完整性提供示例# src/game_manager.pyclass 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.pyfrom fastapi import FastAPIapp = 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.pyasync 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_managerapp.dependency_overrides[get_manager] = override_managerclient = 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应用在各种异常情况下的连接管理逻辑。
以上就是使用PyTest测试FastAPI WebSocket连接的关闭：一种可靠的方法的详细内容，更多请关注创想鸟其它相关文章！
版权声明：本文内容由互联网用户自发贡献，该文观点仅代表作者本人。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。
如发现本站有涉嫌抄袭侵权/违法违规的内容， 请发送邮件至 chuangxiangniao@163.com 举报，一经查实，本站将立刻删除。
发布者：程序猿，转转请注明出处：https://www.chuangxiangniao.com/p/1374748.html
微信扫一扫 
支付宝扫一扫 
