本文详细介绍了在Pytest框架下，如何有效测试FastAPI WebSocket连接的异常关闭情况，特别是当服务器因业务逻辑（如连接无效房间）主动断开连接时。核心解决方案是利用pytest.raises(WebSocketDisconnect)结合客户端尝试接收数据，以捕获服务器断开连接后引发的异常，从而验证连接已按预期关闭。
1. 理解WebSocket连接关闭与测试需求
在开发基于fastapi的websocket应用时，我们经常需要处理各种连接状态，包括正常的建立、通信以及异常的断开。特别是在某些业务逻辑下，服务器可能会主动拒绝或关闭客户端的连接，例如当客户端尝试连接一个不存在的房间或未授权的资源时。为了确保这些断开逻辑的正确性，编写可靠的测试用例至关重要。
使用Pytest作为测试框架，FastAPI提供了TestClient来模拟客户端行为，包括WebSocket连接。我们的目标是验证当服务器主动关闭连接时，客户端能够正确地捕获到WebSocketDisconnect异常。
2. 初始尝试及遇到的问题
最初的测试思路可能是直接将pytest.raises(WebSocketDisconnect)上下文管理器包裹在client.websocket_connect()调用外部，期望连接失败时立即捕获异常。
import pytestfrom fastapi.testclient import TestClientfrom fastapi.websockets import WebSocketDisconnect# 假设 app 和 get_manager, override_manager 等已定义# ... (省略依赖注入和TestClient初始化代码) ...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()在某些情况下可能会成功建立底层的TCP连接，即使服务器端的业务逻辑立即决定关闭WebSocket连接。WebSocketDisconnect异常通常不是在连接建立的瞬间抛出，而是在客户端尝试对一个已经关闭的WebSocket连接执行操作（如接收或发送数据）时才会显现。
在服务器端，当manager.connect()方法在检测到无效房间时抛出WebSocketDisconnect，FastAPI的WebSocket路由会捕获这个异常，并相应地关闭连接。
@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) # 如果 manager.connect 成功，客户端才会被接受并进入消息循环 client_id = websocket.scope["client_id"] while True: data = await websocket.receive_json() await manager.handle_message(room_name, client_id, data) except WebSocketDisconnect: # 当 manager.connect 抛出 WebSocketDisconnect 或客户端断开时，这里会捕获 await manager.remove(websocket) # FastAPI/Starlette 框架会确保连接被正确关闭
3. 正确测试WebSocket连接关闭的方法
解决上述问题的关键在于理解WebSocketDisconnect的触发时机。它通常在客户端尝试从一个已关闭或断开的WebSocket连接接收数据时被抛出。因此，在客户端成功连接（即使服务器立即关闭）后，我们应该主动尝试从连接中接收数据。
3.1 核心思路
客户端使用client.websocket_connect()尝试连接。即使服务器端因业务逻辑立即关闭连接，websocket_connect本身可能不会立即抛出WebSocketDisconnect。在with client.websocket_connect(…) as ws:块内部，立即调用ws.receive_json()（或ws.receive_text()、ws.receive_bytes()）。如果服务器已经关闭了连接，尝试接收数据将触发WebSocketDisconnect异常。将整个过程包裹在pytest.raises(WebSocketDisconnect)中，即可成功捕获并验证异常。
3.2 示例代码
以下是经过修正的、能够正确测试WebSocket连接关闭的Pytest用例：
import pytestfrom fastapi.testclient import TestClientfrom fastapi.websockets import WebSocketDisconnectfrom typing import Annotated# 假设以下导入和定义已存在from src.game_manager import GameManagerfrom src.main import appfrom fastapi import Depends, WebSocket, APIRouter # 假设 router 是一个 APIRouter 实例# 假设 get_manager 是一个依赖函数def get_manager(): # 这是一个占位符，实际应用中会从依赖注入系统获取 GameManager 实例 pass# 重写依赖以进行测试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# 将依赖重写应用于测试客户端app.dependency_overrides[get_manager] = override_manager# 初始化测试客户端client = TestClient(app)class TestWebsocketConnection: def test_connect_to_non_existing_room(self): """ 测试连接到不存在的房间时，WebSocket连接是否被正确关闭。 """ with pytest.raises(WebSocketDisconnect): with client.websocket_connect("/ws/non_existing_room") as ws: # 尝试从已关闭的连接接收数据，这将触发 WebSocketDisconnect ws.receive_json() def test_connect_to_existing_room(self): """ 测试连接到存在的房间时，连接是否成功建立且未断开。 """ with client.websocket_connect("/ws/foo") as ws: # 确保连接成功，可以发送和接收数据 ws.send_json({"message": "Hello from foo"}) # 接收一条消息，验证连接活跃 response = ws.receive_json() assert "type" in response and response["type"] == "websocket.send" # 示例断言 # 确保连接在退出 with 块后正常关闭 with pytest.raises(WebSocketDisconnect): # 尝试在连接关闭后接收数据，验证其已关闭 ws.receive_json()
3.3 注意事项
服务器端行为是关键： 这种测试方法的前提是服务器端确实在某些条件下（如manager.connect失败）主动关闭了WebSocket连接。如果服务器只是不发送数据但保持连接开放，receive_json()将可能无限期等待或超时，而不是抛出WebSocketDisconnect。WebSocketDisconnect的来源： WebSocketDisconnect异常可能由多种情况触发，包括客户端主动断开、服务器主动断开、网络问题等。在测试中，我们确保它是由于服务器端业务逻辑拒绝连接导致的。依赖注入： 在测试中，通常需要通过app.dependency_overrides来模拟或控制依赖项（如GameManager）的行为，以确保测试的隔离性和可预测性。测试覆盖： 除了测试连接关闭，还应测试正常连接、消息发送与接收、以及其他异常情况。
4. 总结
在Pytest中测试FastAPI WebSocket连接的关闭，尤其是当服务器因业务逻辑主动断开连接时，关键在于模拟客户端尝试从一个已关闭的连接接收数据。通过将ws.receive_json()（或类似接收方法）调用包裹在pytest.raises(WebSocketDisconnect)上下文管理器中，我们可以有效地验证服务器是否按预期关闭了连接。这种方法能够精确地捕获到WebSocketDisconnect异常，从而确保WebSocket连接管理逻辑的健壮性。
以上就是在Pytest中测试FastAPI WebSocket连接关闭的正确方法的详细内容，更多请关注创想鸟其它相关文章！
版权声明：本文内容由互联网用户自发贡献，该文观点仅代表作者本人。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。
如发现本站有涉嫌抄袭侵权/违法违规的内容， 请发送邮件至 chuangxiangniao@163.com 举报，一经查实，本站将立刻删除。
发布者：程序猿，转转请注明出处：https://www.chuangxiangniao.com/p/1374864.html
微信扫一扫 
支付宝扫一扫 
