Denny的博客

Denny的博客

0%

FastAPI自学教程（68） - WebSocket测试详解

发表于 分类于

1.基础测试方法

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
from fastapi import FastAPI
from fastapi.testclient import TestClient

app = FastAPI()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    await websocket.send_text("Hello WebSocket")
    await websocket.close()

def test_websocket():
    client = TestClient(app)
    with client.websocket_connect("/ws") as websocket:
        data = websocket.receive_text()
        assert data == "Hello WebSocket"
  • 使用TestClient创建测试客户端，websocket_connect()方法建立WebSocket连接
  • 必须使用with上下文管理器管理连接生命周期，确保正确关闭
  • 流程图展示测试流程： 
    graph TD
  A[启动TestClient] --> B[建立WebSocket连接]
  B --> C[发送模拟请求]
  C --> D[接收响应数据]
  D --> E[断言验证]
  E --> F[关闭连接]

2.数据交互测试方法

2.1 发送数据方法

1
2
3
websocket.send_text("测试消息")  # 发送文本
websocket.send_bytes(b"binary")  # 发送二进制
websocket.send_json({"key": "value"})  # 发送JSON数据
  • 支持三种数据格式发送：文本、二进制、JSON
  • JSON模式可选择mode="binary"通过二进制帧发送

2.2 接收数据方法

1
2
3
text = websocket.receive_text()  # 接收文本
binary = websocket.receive_bytes()  # 接收二进制
json_data = websocket.receive_json()  # 接收JSON
  • 接收方法会阻塞直到数据到达
  • 断言示例： 
    1
    assert json_data == {"status": "success"}

3.异常场景测试

3.1 连接断开测试

1
2
3
4
5
6
7
8
import pytest
from starlette.websockets import WebSocketDisconnect

def test_disconnect():
    with pytest.raises(WebSocketDisconnect):
        with client.websocket_connect("/ws") as websocket:
            websocket.close(code=1001)
            websocket.receive_text()
  • 捕获WebSocketDisconnect异常验证连接关闭
  • 可指定关闭状态码（如1000表示正常关闭）

3.2 错误响应测试

1
2
3
4
def test_invalid_route():
    with pytest.raises(WebSocketDisconnect):
        with client.websocket_connect("/invalid"):
            pass
  • 测试不存在的路由返回404状态码

4.生产环境测试建议

  1. 多客户端模拟： 

    1
    2
    3
    4
    5
    async def multi_client_test():
        async with websockets.connect(uri) as ws1, \
                   websockets.connect(uri) as ws2:
            await ws1.send("msg1")
            await ws2.send("msg2")

    使用websockets库模拟并发连接

  2. 性能基准测试：

    • 通过pytest-benchmark插件测量响应延迟
    • 建议单连接吞吐量>1000 msg/sec

  3. 安全测试要点：

    • 测试最大消息大小限制（默认1MB）
    • 验证wss://协议加密传输
    • 检查跨域请求头Access-Control-Allow-Origin

测试命令示例： 

1
2
3
4
5
# 运行测试用例
pytest test_websockets.py -v

# 带覆盖率检测
pytest --cov=app --cov-report=html