FastAPI 中用 TestClient 测试 WebSocket 连接的完整实践

在构建包含实时通信功能的应用时，WebSocket 端点同样需要自动化测试来保障质量。本篇技术指南以 FastAPI 官方文档中的 Testing WebSockets 章节为核心，讲解如何使用 TestClient 在 with 语句中连接 WebSocket 并断言收发数据；读完本篇，你将掌握在 FastAPI 测试中建立 WebSocket 会话、接收 JSON 数据、验证响应内容的完整套路，并能理解其底层由 Starlette TestClient 直接提供的实现机制。

核心思路：同一个 TestClient 也能测 WebSocket

FastAPI 应用测试使用的 TestClient 并非只能发 HTTP 请求。根据文档 Testing WebSockets，同一个 TestClient 就可以测试 WebSocket，关键技巧在于两点：

1. 把 TestClient 放在一个 with 语句中使用，作为上下文管理器；
2. 在 with 块内调用 client.websocket_connect(...) 建立 WebSocket 连接。

从源码看，fastapi.testclient 实际上是一个"薄封装"——fastapi/testclient.py 的全部内容只有一行：

from starlette.testclient import TestClient as TestClient  # noqa

这意味着 websocket_connect 等 WebSocket 测试能力完整继承自 Starlette 的 TestClient，FastAPI 没有额外改动其行为。文档也明确指出：更多细节（例如连接被拒绝、协议错误等场景）可参考 Starlette 的 testing WebSockets 文档。

完整示例：HTTP 端点与 WebSocket 端点一起测试

下面的示例来自 FastAPI 仓库中的教程源码 docs_src/app_testing/tutorial002_py310.py，它在一个测试文件里同时覆盖了普通 HTTP 端点和 WebSocket 端点，是最典型的写法：

from fastapi import FastAPI
from fastapi.testclient import TestClient
from fastapi.websockets import WebSocket

app = FastAPI()


@app.get("/")
async def read_main():
    return {"msg": "Hello World"}


@app.websocket("/ws")
async def websocket(websocket: WebSocket):
    await websocket.accept()
    await websocket.send_json({"msg": "Hello WebSocket"})
    await websocket.close()


def test_read_main():
    client = TestClient(app)
    response = client.get("/")
    assert response.status_code == 200
    assert response.json() == {"msg": "Hello World"}


def test_websocket():
    client = TestClient(app)
    with client.websocket_connect("/ws") as websocket:
        data = websocket.receive_json()
        assert data == {"msg": "Hello WebSocket"}

示例拆解

服务端部分：@app.websocket("/ws") 注册了一个 WebSocket 路由。注意三个方法调用缺一不可：

• await websocket.accept()：必须先接受连接，否则客户端会收到连接失败；
• await websocket.send_json(...)：发送 JSON 文本帧；
• await websocket.close()：服务端主动关闭连接。

HTTP 测试部分：test_read_main 展示了常规用法——TestClient 不必放在 with 语句里，普通的 client.get("/") 即可断言状态码和 JSON 响应体。

WebSocket 测试部分：test_websocket 是本文档的核心，逐行看：

def test_websocket():
    client = TestClient(app)
    with client.websocket_connect("/ws") as websocket:   # 1. 用 with 语句建立连接
        data = websocket.receive_json()                  # 2. 接收 JSON 帧
        assert data == {"msg": "Hello WebSocket"}         # 3. 断言收到的数据

1. with client.websocket_connect("/ws")：websocket_connect 返回的上下文管理器负责完成 WebSocket 握手（HTTP Upgrade 请求）与连接生命周期管理。进入 with 块时连接已建立，离开 with 块时连接自动关闭。路径 "/ws" 就是 @app.websocket 装饰器声明的路由。
2. websocket.receive_json()：接收服务端发送的第一个 JSON 消息，并自动解析为 Python 对象（这里是 dict）。文档中示例高亮的正是这一段（第 27–31 行）。
3. assert：直接对解析后的 Python 对象做等值断言，无需手动解析字符串。

为什么必须用 with 语句？

websocket_connect 返回的是上下文管理器，with 语句保证连接在测试结束后被正确清理。如果不使用 with 而是长期持有连接，测试之间可能产生资源泄漏，且无法利用其退出时的自动关闭逻辑。仓库中其他 WebSocket 测试也遵循同样的模式，例如 tests/test_ws_dependencies.py 中的写法：

def test_index():
    client = TestClient(app)
    with client.websocket_connect("/") as websocket:
        data = json.loads(websocket.receive_text())
        assert data == ["app", "index"]

可以看到该测试使用 receive_text() 接收原始文本再手动 json.loads，与教程中直接 receive_json() 等价，选择哪种取决于服务端发送的是 send_text 还是 send_json。

WebSocket 测试中可用的收发 API

在上述模式基础上，websocket 对象在 with 块内提供了一组对称的收发方法，覆盖 WebSocket 协议的常见帧类型：

方法	作用
websocket.receive_json()	接收 JSON 帧并解析为 Python 对象
websocket.receive_text()	接收文本帧，返回原始字符串
websocket.receive_bytes()	接收二进制帧
websocket.send_json(obj)	向服务端发送 JSON 帧
websocket.send_text(text)	向服务端发送文本帧
websocket.send_bytes(data)	向服务端发送二进制帧
websocket.close(code=1000)	客户端主动关闭连接，可携带关闭状态码

这些方法在仓库测试中有真实使用佐证。例如 tests/test_dependency_after_yield_websockets.py、tests/test_ws_dependencies.py、tests/test_route_scope.py 均通过 client.websocket_connect(...) 建立会话并断言往返数据，验证了 WebSocket 端点在依赖注入、路由前缀、Router 挂载等场景下的行为。

与 FastAPI WebSocket API 的配合

测试对象是服务端路由中使用的 WebSocket 类型。从 fastapi/websockets.py 可以确认，WebSocket 相关类型同样是复用 Starlette 的实现：

from starlette.websockets import WebSocketDisconnect as WebSocketDisconnect  # noqa

因此在编写 WebSocket 端点时，异常处理依然使用熟悉的 WebSocketDisconnect 捕获客户端意外断开；而测试侧的 websocket_connect 与服务端 websocket.accept()、websocket.send_json() 等方法天然对接，测试代码与服务端代码使用同一套语义。

运行与验证方式

将上面的示例保存为一个 pytest 可收集的测试文件后，在项目环境中执行：

pytest -s

预期 test_read_main 与 test_websocket 均通过。若断言失败（例如服务端忘记 accept() 或发送的数据与预期不符），pytest 会报出 WebSocketDisconnect 或断言不等的具体信息，便于定位问题。

仓库自身也使用相同的验证脚本，如 scripts/test.sh，可参考其依赖与执行方式来配置本地测试环境。

小结

• FastAPI 的 TestClient 直接复用自 Starlette，with client.websocket_connect(path) 是测试 WebSocket 的标准入口；
• 测试代码中通过 receive_json() / receive_text() 等方法接收数据并直接断言 Python 对象，无需手工解析；
• 同一测试文件可以同时覆盖 HTTP 端点与 WebSocket 端点，完整示例见 docs_src/app_testing/tutorial002_py310.py；
• 更复杂的连接失败、协议错误等边界场景，其行为以 Starlette 的 TestClient 实现为准。