veza/tools/tests/perf/k6_chat_ws.js

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';

// Load environment variables
const API_ORIGIN = __ENV.API_ORIGIN || 'https://api.lab.veza';
const CHAT_ORIGIN = __ENV.CHAT_ORIGIN || 'wss://chat.lab.veza';
const TEST_EMAIL_PREFIX = __ENV.TEST_EMAIL_PREFIX || 'user+ws';
const TEST_EMAIL_DOMAIN = __ENV.TEST_EMAIL_DOMAIN || 'lab.veza';
const TEST_PASSWORD_PREFIX = __ENV.TEST_PASSWORD_PREFIX || 'V3za!ws-';

// Custom metrics
const wsConnectionTime = new Trend('ws_connection_time');
const wsMessageSendTime = new Trend('ws_message_send_time');
const wsMessageReceiveTime = new Trend('ws_message_receive_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');

// Test configuration
export const options = {
  scenarios: {
    websocket_chat: {
      executor: 'ramping-vus',
      stages: [
        { duration: '30s', target: 10 },  // Ramp up to 10 connections
        { duration: '2m', target: 50 },   // Ramp up to 50 connections
        { duration: '2m', target: 50 },   // Stay at 50 connections
        { duration: '30s', target: 0 },   // Ramp down
      ],
      gracefulRampDown: '30s',
    },
  },
  thresholds: {
    ws_connection_time: ['p(95)<500', 'p(99)<1000'],
    ws_message_send_time: ['p(95)<100', 'p(99)<200'],
    ws_connection_failures: ['rate<0.05'], // Less than 5% connection failures
    ws_message_failures: ['rate<0.01'],    // Less than 1% message failures
  },
};

// Helper to create and authenticate a user
function createAuthenticatedUser() {
  const user = {
    email: `${TEST_EMAIL_PREFIX}${randomString(8)}@${TEST_EMAIL_DOMAIN}`,
    password: `${TEST_PASSWORD_PREFIX}${randomString(8)}`,
  };
  
  // Register user
  const registerRes = http.post(`${API_ORIGIN}/auth/register`, JSON.stringify(user), {
    headers: { 'Content-Type': 'application/json' },
  });
  
  if (registerRes.status !== 201) {
    console.error(`Failed to register user: ${registerRes.status}`);
    return null;
  }
  
  try {
    const body = JSON.parse(registerRes.body);
    return {
      ...user,
      accessToken: body.access_token,
      userId: body.user.id,
    };
  } catch (e) {
    console.error('Failed to parse registration response');
    return null;
  }
}

// Setup: Create test users
export function setup() {
  const users = [];
  const userCount = 60; // Create more than we need
  
  console.log(`Creating ${userCount} test users for WebSocket test...`);
  
  for (let i = 0; i < userCount; i++) {
    const user = createAuthenticatedUser();
    if (user) {
      users.push(user);
    }
    sleep(0.1); // Small delay to avoid overwhelming the API
  }
  
  console.log(`Created ${users.length} authenticated users`);
  return { users };
}

export default function (data) {
  const { users } = data;
  
  // Pick a random user
  const user = users[Math.floor(Math.random() * users.length)];
  if (!user) {
    console.error('No user available');
    return;
  }
  
  const url = `${CHAT_ORIGIN}?token=${user.accessToken}`;
  const params = { tags: { user: user.email } };
  
  let connectionStartTime = new Date();
  let connected = false;
  let messageCount = 0;
  let lastMessageTime = new Date();
  
  const res = ws.connect(url, params, function (socket) {
    socket.on('open', () => {
      const connectionEndTime = new Date();
      wsConnectionTime.add(connectionEndTime - connectionStartTime);
      connected = true;
      
      console.log(`User ${user.email} connected to WebSocket`);
      
      // Join general room
      socket.send(JSON.stringify({
        type: 'join',
        room: 'general',
      }));
      
      // Send periodic messages
      socket.setInterval(() => {
        if (connected) {
          const message = {
            type: 'message',
            room: 'general',
            content: `Test message from ${user.email} at ${new Date().toISOString()}`,
          };
          
          const sendStart = new Date();
          socket.send(JSON.stringify(message));
          messagesSent.add(1);
          
          // Record send time (approximation)
          wsMessageSendTime.add(new Date() - sendStart);
        }
      }, 5000); // Send message every 5 seconds
      
      // Send ping every 30 seconds to keep connection alive
      socket.setInterval(() => {
        if (connected) {
          socket.send(JSON.stringify({ type: 'ping' }));
        }
      }, 30000);
    });
    
    socket.on('message', (data) => {
      const receiveTime = new Date();
      messagesReceived.add(1);
      messageCount++;
      
      try {
        const message = JSON.parse(data);
        
        // Calculate receive time since last message
        wsMessageReceiveTime.add(receiveTime - lastMessageTime);
        lastMessageTime = receiveTime;
        
        // Check message validity
        check(message, {
          'message has type': (m) => m.type !== undefined,
          'message is valid': (m) => m.error === undefined,
        });
        
        // Handle different message types
        switch (message.type) {
          case 'joined':
            console.log(`User ${user.email} joined room`);
            break;
          case 'message':
            // Received a chat message
            break;
          case 'error':
            console.error(`WebSocket error for ${user.email}: ${message.message}`);
            wsMessageFailures.add(1);
            break;
        }
      } catch (e) {
        console.error(`Failed to parse message: ${data}`);
        wsMessageFailures.add(1);
      }
    });
    
    socket.on('close', () => {
      connected = false;
      console.log(`User ${user.email} disconnected. Received ${messageCount} messages`);
    });
    
    socket.on('error', (e) => {
      console.error(`WebSocket error for ${user.email}: ${e}`);
      wsConnectionFailures.add(1);
    });
    
    // Simulate user activity for 30-60 seconds
    const activityDuration = 30 + Math.random() * 30;
    socket.setTimeout(() => {
      console.log(`User ${user.email} ending session after ${activityDuration}s`);
      socket.close();
    }, activityDuration * 1000);
  });
  
  // Check connection result
  check(res, {
    'WebSocket connection successful': (r) => r && r.status === 101,
  });
  
  if (!res || res.status !== 101) {
    wsConnectionFailures.add(1);
  }
}

// Teardown
export function teardown(data) {
  console.log('WebSocket chat test completed');
  console.log(`Total users created: ${data.users.length}`);
}