如何使用aiohttp编写含Class和task的Locust WebSocket测试脚本

以下是符合需求的Locust脚本，采用异步模式结合aiohttp处理WebSocket连接，完全遵循Locust的类和task特性：

from locust import AsyncUser, task, run_single_user
from aiohttp import WSMsgType

class WebSocketUser(AsyncUser):
    # 配置WebSocket服务的基础地址
    host = "wss://someurl"

    @task
    async def websocket_interaction(self):
        # 建立WebSocket连接，Locust会自动统计该请求的状态指标
        async with self.client.ws_connect("/api/task/ws") as ws:
            # 示例：发送一条JSON格式消息到服务端（根据实际业务调整内容）
            await ws.send_json({"action": "ping", "data": "test_payload"})
            
            # 示例：循环接收服务端返回的消息
            async for msg in ws:
                if msg.type == WSMsgType.TEXT:
                    # 处理文本消息，可根据业务逻辑添加断言或日志
                    print(f"收到消息: {msg.data}")
                    # 收到消息后主动关闭连接（可根据需求改为保持长连接）
                    await ws.close()
                    break
                elif msg.type == WSMsgType.ERROR:
                    # 捕获WebSocket错误
                    print(f"WebSocket错误: {ws.exception()}")
                    break

# 单用户调试入口（可选）
if __name__ == "__main__":
    run_single_user(WebSocketUser)

关键说明： #

• 使用AsyncUser作为基类：适配aiohttp的异步模型，支持高并发WebSocket连接
• ws_connect方法：Locust异步客户端内置的WebSocket连接方法，会自动校验101 Switching Protocols状态码，同时将连接请求纳入Locust的性能统计
• @task装饰器：标记该方法为压测任务，Locust会根据配置的并发用户数循环执行该任务
• 消息交互逻辑：示例包含基础的收发消息逻辑，可根据实际业务场景修改（比如发送二进制消息、保持长连接持续交互等）

运行方式： #

1. 保存脚本为websocket_locust.py
2. 执行命令启动压测：locust -f websocket_locust.py
3. 在浏览器打开http://localhost:8089，配置并发用户数、孵化率等参数后开始压测

内容的提问来源于stack exchange，提问作者Magnus Jensen