本文探讨了如何使用PyTest有效测试FastAPI WebSocket连接在服务器端被主动关闭的场景。当服务器因业务逻辑（如连接到不存在的房间）而拒绝或关闭连接时，客户端尝试接收数据会触发WebSocketDisconnect异常。通过在PyTest中捕获此异常，可以验证连接关闭的预期行为。
背景与挑战
在开发基于fastapi的websocket应用程序时，尤其是在构建需要严格连接管理的游戏或其他实时通信服务时，测试服务器端主动关闭连接的场景至关重要。例如，当客户端尝试连接到一个不存在的房间时，服务器理应立即拒绝并关闭该连接。
然而，在使用PyTest和fastapi.testclient.TestClient进行测试时，直接在client.websocket_connect上下文管理器外部或内部尝试捕获WebSocketDisconnect异常，往往会发现测试失败，并提示“DID NOT RAISE ”。这表明WebSocketDisconnect异常并非在websocket_connect调用时立即抛出，而是发生在后续的某个时刻。
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 pytestfrom fastapi.testclient import TestClientfrom 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 pytestfrom fastapi.testclient import TestClientfrom fastapi.websockets import WebSocketDisconnectfrom typing import Annotated# 假设你的FastAPI应用和GameManager都在src目录下from src.game_manager import GameManagerfrom 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应用程序的健壮性提供了可靠的测试保障。
以上就是PyTest测试FastAPI WebSocket连接关闭的有效策略的详细内容，更多请关注创想鸟其它相关文章！
版权声明：本文内容由互联网用户自发贡献，该文观点仅代表作者本人。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。
如发现本站有涉嫌抄袭侵权/违法违规的内容， 请发送邮件至 chuangxiangniao@163.com 举报，一经查实，本站将立刻删除。
发布者：程序猿，转转请注明出处：https://www.chuangxiangniao.com/p/1374918.html
微信扫一扫 
支付宝扫一扫 
