PyTest测试FastAPI WebSocket连接关闭的有效策略-Python教程-PHP中文网

PyTest测试FastAPI WebSocket连接关闭的有效策略

DDD

发布： 2025-10-01 15:46:29

原创

875人浏览过

PyTest测试FastAPI WebSocket连接关闭的有效策略

本文探讨了如何使用PyTest有效测试FastAPI WebSocket连接在服务器端被主动关闭的场景。当服务器因业务逻辑（如连接到不存在的房间）而拒绝或关闭连接时，客户端尝试接收数据会触发WebSocketDisconnect异常。通过在PyTest中捕获此异常，可以验证连接关闭的预期行为。

背景与挑战

在开发基于fastapi的websocket应用程序时，尤其是在构建需要严格连接管理的游戏或其他实时通信服务时，测试服务器端主动关闭连接的场景至关重要。例如，当客户端尝试连接到一个不存在的房间时，服务器理应立即拒绝并关闭该连接。

然而，在使用PyTest和fastapi.testclient.TestClient进行测试时，直接在client.websocket_connect上下文管理器外部或内部尝试捕获WebSocketDisconnect异常，往往会发现测试失败，并提示“DID NOT RAISE <class 'starlette.websockets.WebSocketDisconnect'>”。这表明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)：

白瓜面试

白瓜面试 - AI面试助手,辅助笔试面试神器

查看详情

# 错误的测试方法示例
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应用程序的健壮性提供了可靠的测试保障。

以上就是PyTest测试FastAPI WebSocket连接关闭的有效策略的详细内容，更多请关注php中文网其它相关文章！