veza/loadtests/chat/websocket.js

/**
 * Load test: WebSocket chat
 * Usage: k6 run loadtests/chat/websocket.js
 * Requires: Backend API + Chat server running
 * Flow: register -> login -> GET /api/v1/chat/token -> connect WS with token
 */
import ws from 'k6/ws';
import { check, sleep } from 'k6';
import { Rate, Trend, Counter } from 'k6/metrics';
import { randomString } from 'https://jslib.k6.io/k6-utils/1.2.0/index.js';
import http from 'k6/http';

const API_ORIGIN = __ENV.API_ORIGIN || __ENV.BASE_URL || 'http://localhost:8080';
const CHAT_ORIGIN = __ENV.CHAT_ORIGIN || 'ws://localhost:8080';
// Chat migrated to Go backend (v0.502) — WebSocket at /api/v1/ws
const WS_BASE = CHAT_ORIGIN.startsWith('ws') ? CHAT_ORIGIN : `ws://${CHAT_ORIGIN.replace(/^https?:\/\//, '')}`;
const WS_URL = `${WS_BASE.replace(/\/?$/, '')}/api/v1/ws`;
const TEST_EMAIL_PREFIX = __ENV.TEST_EMAIL_PREFIX || 'user+ws';
const TEST_EMAIL_DOMAIN = __ENV.TEST_EMAIL_DOMAIN || 'example.com';
const TEST_PASSWORD_PREFIX = __ENV.TEST_PASSWORD_PREFIX || 'V3za!ws-';

const wsConnectionTime = new Trend('ws_connection_time');
const wsConnectionFailures = new Rate('ws_connection_failures');
const wsMessageFailures = new Rate('ws_message_failures');
const messagesReceived = new Counter('messages_received');
const messagesSent = new Counter('messages_sent');

export const options = {
  scenarios: {
    websocket_chat: {
      executor: 'ramping-vus',
      stages: [
        { duration: '30s', target: 10 },
        { duration: '2m', target: 30 },
        { duration: '1m', target: 30 },
        { duration: '30s', target: 0 },
      ],
      gracefulRampDown: '30s',
    },
  },
  thresholds: {
    ws_connection_time: ['p(95)<500', 'p(99)<1000'],
    ws_connection_failures: ['rate<0.05'],
    ws_message_failures: ['rate<0.05'],
  },
};

function createAuthenticatedUser() {
  const rand = randomString(8);
  const pwd = `${TEST_PASSWORD_PREFIX}${rand}`;
  const user = {
    email: `${TEST_EMAIL_PREFIX}${rand}@${TEST_EMAIL_DOMAIN}`,
    password: pwd,
    password_confirmation: pwd,
    username: `ws${rand}`,
  };

  const registerRes = http.post(`${API_ORIGIN}/api/v1/auth/register`, JSON.stringify(user), {
    headers: { 'Content-Type': 'application/json' },
  });
  if (registerRes.status !== 201) return null;

  const loginRes = http.post(`${API_ORIGIN}/api/v1/auth/login`, JSON.stringify({ email: user.email, password: user.password }), {
    headers: { 'Content-Type': 'application/json' },
  });
  if (loginRes.status !== 200) return null;

  let accessToken = '';
  try {
    const body = JSON.parse(loginRes.body);
    accessToken = body.data?.token?.access_token || '';
  } catch (e) {
    return null;
  }
  if (!accessToken) return null;

  const chatTokenRes = http.get(`${API_ORIGIN}/api/v1/chat/token`, {
    headers: { Authorization: `Bearer ${accessToken}` },
  });
  if (chatTokenRes.status !== 200) return null;

  let chatToken = '';
  try {
    const ctBody = JSON.parse(chatTokenRes.body);
    chatToken = ctBody.data?.token || ctBody.token || '';
  } catch (e) {
    return null;
  }
  if (!chatToken) return null;

  return { ...user, chatToken };
}

export function setup() {
  const users = [];
  for (let i = 0; i < 40; i++) {
    const user = createAuthenticatedUser();
    if (user) users.push(user);
    sleep(0.1);
  }
  return { users };
}

export default function (data) {
  const { users } = data;
  if (users.length === 0) return;

  const user = users[Math.floor(Math.random() * users.length)];
  const url = `${WS_URL}?token=${user.chatToken}`;
  const connectionStart = Date.now();

  const res = ws.connect(url, {}, function (socket) {
    socket.on('open', () => {
      wsConnectionTime.add(Date.now() - connectionStart);

      socket.send(JSON.stringify({ type: 'join', room: 'general' }));

      socket.setInterval(() => {
        socket.send(JSON.stringify({
          type: 'message',
          room: 'general',
          content: `Load test ${Date.now()}`,
        }));
        messagesSent.add(1);
      }, 5000);

      socket.setTimeout(() => {
        socket.close();
      }, 20000 + Math.random() * 10000);
    });

    socket.on('message', (data) => {
      messagesReceived.add(1);
      try {
        const msg = JSON.parse(data);
        if (msg.error) wsMessageFailures.add(1);
      } catch (e) {
        wsMessageFailures.add(1);
      }
    });

    socket.on('error', () => {
      wsConnectionFailures.add(1);
    });
  });

  check(res, { 'WebSocket connection successful': (r) => r && r.status === 101 });
  if (!res || res.status !== 101) {
    wsConnectionFailures.add(1);
  }
}