在FastAPI应用中，使用Pytest测试WebSocket连接是否按预期关闭，尤其是在服务器端因业务逻辑立即拒绝连接时，是一个常见的挑战。本文将深入探讨如何通过尝试从已关闭的WebSocket连接接收数据来触发WebSocketDisconnect异常，从而有效地验证连接的正确关闭，并提供详细的测试代码示例和专业解析。
1. 理解WebSocket连接关闭测试的挑战
在开发基于fastapi的实时应用（如在线游戏）时，websocket通信是核心。测试这些应用的健壮性，包括在特定条件下连接应立即关闭的场景，至关重要。例如，当客户端尝试连接到一个不存在的房间或未授权的资源时，服务器应立即终止连接。
使用Pytest对FastAPI的WebSocket端点进行测试时，一个常见的误区是直接将pytest.raises(WebSocketDisconnect)应用于client.websocket_connect上下文管理器。原始尝试如下：
import pytestfrom fastapi.testclient import TestClientfrom fastapi.websockets import WebSocketDisconnect# ... (省略导入和应用初始化代码)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 。这是因为client.websocket_connect本身主要负责完成WebSocket握手过程。即使服务器在握手完成后立即关闭连接，websocket_connect上下文管理器本身通常不会立即抛出WebSocketDisconnect。WebSocketDisconnect异常通常在尝试对一个已经关闭的连接进行读写操作时才会被触发。
2. FastAPI WebSocket端点的工作原理
为了更好地理解测试方法，我们回顾一下FastAPI WebSocket端点的基本结构：
from fastapi import APIRouter, WebSocket, Dependsfrom fastapi.websockets import WebSocketDisconnectfrom typing import Annotated# 假设 GameManager 和 get_manager 已定义from src.game_manager import GameManager, get_manager 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: # manager.connect 负责处理连接逻辑，如果条件不满足，它会抛出 WebSocketDisconnect await manager.connect(websocket, room_name, password) # 获取客户端ID 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)
关键点在于manager.connect(websocket, room_name, password)方法。根据描述，如果连接不被允许（例如，房间不存在），此方法会立即抛出WebSocketDisconnect。当manager.connect抛出此异常时，websocket_endpoint中的except WebSocketDisconnect块会被执行，进而调用manager.remove(websocket)，并在内部可能调用websocket.close()来显式关闭服务器端的WebSocket连接。
3. 正确测试WebSocket连接关闭的方法
既然WebSocketDisconnect是在对已关闭连接进行操作时抛出的，那么解决方案就是：在连接建立后，立即尝试从客户端接收数据。如果服务器确实在没有发送任何数据的情况下关闭了连接，客户端的ws.receive_json()操作将立即感知到连接的关闭，并抛出WebSocketDisconnect。
以下是修正后的测试代码：
import pytestfrom fastapi.testclient import TestClientfrom fastapi.websockets import WebSocketDisconnect# 假设 src.game_manager 和 src.main 已正确导入from src.game_manager import GameManager, get_managerfrom src.main import app# 依赖覆盖，用于测试环境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.managerapp.dependency_overrides[get_manager] = override_managerclient = TestClient(app)class TestWebsocketConnection: def test_connect_to_non_existing_room_correct(self): # 期望 WebSocketDisconnect 异常被抛出 with pytest.raises(WebSocketDisconnect): # 建立 WebSocket 连接 with client.websocket_connect("/ws/non_existing_room") as ws: # 尝试接收数据。由于服务器会立即关闭连接，此操作会抛出 WebSocketDisconnect ws.receive_json()
解析：
with client.websocket_connect(“/ws/non_existing_room”) as ws:：这一步会尝试与FastAPI应用建立WebSocket连接。即使服务器端因业务逻辑（如manager.connect抛出异常）立即关闭连接，客户端的websocket_connect上下文管理器本身通常仍能完成初始的握手过程。ws.receive_json()：在client.websocket_connect上下文内部，我们立即尝试从服务器接收JSON数据。如果服务器端在manager.connect中抛出WebSocketDisconnect，并因此关闭了WebSocket连接，那么客户端尝试执行ws.receive_json()时，会发现连接已关闭，从而在客户端侧也抛出WebSocketDisconnect异常。with pytest.raises(WebSocketDisconnect):：这个Pytest结构会捕获并验证内部代码块是否抛出了WebSocketDisconnect异常。如果抛出，测试通过；否则，测试失败。
4. 注意事项与最佳实践
异常时机： WebSocketDisconnect异常通常在尝试对一个已关闭的WebSocket连接进行读写操作时发生。它不是连接建立失败的通用指示器，而是连接“意外”断开或服务器主动关闭后客户端尝试交互的信号。服务器行为： 确保服务器端的业务逻辑确实会在不满足条件时调用websocket.close()或通过抛出WebSocketDisconnect来触发连接关闭。如果服务器只是简单地不发送数据但不关闭连接，客户端的receive_json()可能会无限期等待。测试粒度： 这种测试方法适用于验证服务器端在特定业务逻辑下立即拒绝并关闭连接的场景。对于更复杂的连接生命周期测试（如客户端主动断开、服务器超时断开等），可能需要结合更精细的模拟和异步测试技术。依赖注入： 在测试中，通过app.dependency_overrides替换get_manager依赖，允许我们为测试环境配置一个可控的GameManager实例，确保测试的隔离性和可重复性。
5. 总结
通过在pytest.raises(WebSocketDisconnect)块内尝试从客户端的WebSocket连接接收数据，我们能够可靠地测试FastAPI应用中WebSocket连接是否按预期关闭。这种方法利用了WebSocketDisconnect异常的触发机制，即在尝试与已关闭连接交互时抛出，从而有效地验证了服务器端立即拒绝并终止连接的业务逻辑。掌握这一技巧，对于构建健壮且可测试的FastAPI WebSocket应用至关重要。
以上就是FastAPI WebSocket连接关闭测试：Pytest实战指南的详细内容，更多请关注创想鸟其它相关文章！
版权声明：本文内容由互联网用户自发贡献，该文观点仅代表作者本人。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。
如发现本站有涉嫌抄袭侵权/违法违规的内容， 请发送邮件至 chuangxiangniao@163.com 举报，一经查实，本站将立刻删除。
发布者：程序猿，转转请注明出处：https://www.chuangxiangniao.com/p/1374860.html
微信扫一扫 
支付宝扫一扫 
