Artillery 中使用自定义处理器为 WebSocket 用户生成 JWT 令牌

在 Artillery 负载测试工具中，当我们需要为每个虚拟用户(VU)生成不同的 JWT 令牌来测试 WebSocket 服务时，正确的配置方式至关重要。本文将详细介绍如何实现这一需求。

问题背景

在测试需要认证的 WebSocket 服务时，通常需要在连接时传递 JWT 令牌。每个测试用户应该使用不同的令牌，以模拟真实场景。常见的错误做法是直接在配置层面调用自定义函数，这会导致令牌无法正确生成和传递。

正确实现方式

1. 创建处理器文件

首先创建一个 processor.js 文件，其中包含生成随机令牌的逻辑：

// 预定义的令牌数组
let tokens = [
  "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9...",
  "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9...",
  "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9..."
];

// 生成随机令牌的函数
function generateRandomToken(userContext, events, done) {
  // 从数组中随机选择一个令牌
  userContext.vars.token = tokens[Math.floor(Math.random() * tokens.length)];
  return done();
}

module.exports = {
    generateRandomToken
};

2. 编写测试脚本

在测试脚本中，关键点是将函数调用放在场景流程中，而不是配置层面：

config:
  target: "ws://localhost:8080/userMessages"
  phases:
    - duration: 60  # 测试持续时间(秒)
      arrivalRate: 10  # 每秒创建的用户数
  processor: "./processor.js"  # 指定处理器文件路径
  ws:
    headers:
      x-jwt-token: "{{ token }}"  # 使用模板变量引用令牌

scenarios:
  - name: "WebSocket 认证测试"
    engine: "ws"
    flow:
      - function: "generateRandomToken"  # 在流程中调用函数
      - send: "Test message"  # 发送测试消息
      - think: 5  # 思考时间(秒)

技术要点解析

1. 处理器函数调用位置：必须在场景流程中调用，确保每个虚拟用户都能获得独立的令牌值。

2. 变量作用域：userContext.vars 中设置的变量只在当前虚拟用户会话中有效。

3. 令牌生成策略：示例中使用预定义数组随机选择，实际应用中可根据需要：

   • 动态生成有效JWT
   • 从外部文件读取
   • 调用认证接口获取

4. WebSocket 头设置：通过 ws.headers 配置认证头，使用模板语法 {{ token }} 引用变量。

进阶用法

1. 动态令牌生成：可以修改处理器函数，使用库如 jsonwebtoken 动态生成有效JWT：

const jwt = require('jsonwebtoken');

function generateDynamicToken(userContext, events, done) {
  const payload = { userId: userContext.vars.$uuid };
  const token = jwt.sign(payload, 'your-secret-key', { expiresIn: '1h' });
  userContext.vars.token = token;
  return done();
}

2. 令牌生命周期管理：对于长时间运行的测试，可以在处理器中添加令牌刷新逻辑。

3. 性能考虑：对于大规模测试，预生成令牌比动态生成更高效。

常见问题解决

1. 令牌未定义错误：确保函数在发送请求前被调用，且变量名拼写一致。

2. 令牌过期问题：适当设置JWT有效期，或实现刷新机制。

3. 性能瓶颈：避免在处理器中进行复杂计算，必要时预生成测试数据。

通过以上方法，可以有效地为 Artillery 负载测试中的每个 WebSocket 连接生成和使用独立的 JWT 令牌，从而更真实地模拟生产环境中的用户行为。