ShieldAI/load-tests/darkwatch-auth/darkwatch-auth.js

import http from 'k6/http';
import { check, sleep } from 'k6';
import { Rate, Trend } from 'k6/metrics';

// ── Configuration ────────────────────────────────────────────────────────────
const BASE_URL = __ENV.DARKWATCH_BASE_URL || 'http://localhost:3000';
const TEST_EMAIL = __ENV.TEST_EMAIL || 'loadtest@darkwatch.shieldai';
const TEST_PASSWORD = __ENV.TEST_PASSWORD || 'LoadTest2026!';
const DURATION = __ENV.DURATION || '300s'; // 5 minutes
const TARGET_RPS = parseInt(__ENV.TARGET_RPS || '500', 10);

// P99 latency thresholds (ms)
const THRESHOLDS = {
  login: parseInt(__ENV.LOGIN_P99_MS || '200', 10),
  logout: parseInt(__ENV.LOGOUT_P99_MS || '100', 10),
  refresh: parseInt(__ENV.REFRESH_P99_MS || '150', 10),
};

// ── Custom Metrics ───────────────────────────────────────────────────────────
const loginLatency = new Trend('login_p99');
const logoutLatency = new Trend('logout_p99');
const refreshLatency = new Trend('refresh_p99');

const loginSuccess = new Rate('login_success');
const logoutSuccess = new Rate('logout_success');
const refreshSuccess = new Rate('refresh_success');

// ── Helpers ──────────────────────────────────────────────────────────────────
function uuidv4() {
  return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, (c) => {
    const r = (Math.random() * 16) | 0;
    const v = c === 'x' ? r : (r & 0x3) | 0x8;
    return v.toString(16);
  });
}

const authHeaders = {
  'Content-Type': 'application/json',
};

// ── Scenario: Login (POST /auth/login) ──────────────────────────────────────
function testLogin() {
  const uniqueEmail = `${TEST_EMAIL.replace('@', `_${uuidv4().slice(0, 8)}@`)}`;

  const payload = JSON.stringify({
    email: uniqueEmail,
    password: TEST_PASSWORD,
  });

  const res = http.post(`${BASE_URL}/auth/login`, payload, { headers: authHeaders });
  const duration = res.timings.duration;
  loginLatency.add(duration);

  const success = res.status === 200 || res.status === 201;
  loginSuccess.add(success);

  check(res, {
    'login: status 200 or 201': (r) => r.status === 200 || r.status === 201,
    'login: has access_token': (r) => {
      try {
        const json = JSON.parse(r.body);
        return !!json.access_token || !!json.token || !!json.data?.access_token;
      } catch {
        return false;
      }
    },
    `login: P99 < ${THRESHOLDS.login}ms`: (r) => duration < THRESHOLDS.login,
  });

  try {
    const json = JSON.parse(res.body);
    return {
      accessToken: json.access_token || json.token || json.data?.access_token || uuidv4(),
      refreshToken: json.refresh_token || json.data?.refresh_token || uuidv4(),
      userId: json.user?.id || json.data?.user?.id || uuidv4(),
    };
  } catch {
    return {
      accessToken: uuidv4(),
      refreshToken: uuidv4(),
      userId: uuidv4(),
    };
  }
}

// ── Scenario: Refresh (POST /auth/refresh) ──────────────────────────────────
function testRefresh(refreshToken) {
  const token = refreshToken || uuidv4();

  const payload = JSON.stringify({
    refresh_token: token,
  });

  const res = http.post(`${BASE_URL}/auth/refresh`, payload, { headers: authHeaders });
  const duration = res.timings.duration;
  refreshLatency.add(duration);

  const success = res.status === 200;
  refreshSuccess.add(success);

  check(res, {
    'refresh: status 200': (r) => r.status === 200,
    'refresh: has new access_token': (r) => {
      try {
        const json = JSON.parse(r.body);
        return !!json.access_token || !!json.token || !!json.data?.access_token;
      } catch {
        return false;
      }
    },
    `refresh: P99 < ${THRESHOLDS.refresh}ms`: (r) => duration < THRESHOLDS.refresh,
  });

  try {
    const json = JSON.parse(res.body);
    return {
      accessToken: json.access_token || json.token || json.data?.access_token || uuidv4(),
      refreshToken: json.refresh_token || json.data?.refresh_token || token,
    };
  } catch {
    return {
      accessToken: uuidv4(),
      refreshToken: token,
    };
  }
}

// ── Scenario: Logout (POST /auth/logout) ────────────────────────────────────
function testLogout(accessToken) {
  const token = accessToken || uuidv4();

  const payload = JSON.stringify({
    access_token: token,
  });

  const res = http.post(`${BASE_URL}/auth/logout`, payload, {
    headers: {
      ...authHeaders,
      Authorization: `Bearer ${token}`,
    },
  });
  const duration = res.timings.duration;
  logoutLatency.add(duration);

  const success = res.status === 200 || res.status === 204;
  logoutSuccess.add(success);

  check(res, {
    'logout: status 200 or 204': (r) => r.status === 200 || r.status === 204,
    `logout: P99 < ${THRESHOLDS.logout}ms`: (r) => duration < THRESHOLDS.logout,
  });
}

// ── Default Scenario: Weighted mixed workload ────────────────────────────────
export const options = {
  scenarios: {
    sustained_load: {
      executor: 'constant-arrival-rate',
      duration: DURATION,
      rate: TARGET_RPS,
      preAllocatedVUs: 20,
      maxVUs: 100,
      startTime: '0s',
      exec: 'mixedWorkload',
      tags: { scenario: 'sustained_load' },
    },
  },
  thresholds: {
    `login_p99`: [`p(99)<${THRESHOLDS.login}`],
    `logout_p99`: [`p(99)<${THRESHOLDS.logout}`],
    `refresh_p99`: [`p(99)<${THRESHOLDS.refresh}`],
    `login_success`: ['rate>0.95'],
    `logout_success`: ['rate>0.95'],
    `refresh_success`: ['rate>0.95'],
    http_req_duration: [`p(95)<300`, `p(99)<400`],
    http_req_failed: ['rate<0.05'],
  },
};

// Mixed workload: 40% login, 35% refresh, 25% logout
export function mixedWorkload() {
  const rand = Math.random();

  if (rand < 0.4) {
    const tokens = testLogin();
    sleep(0.05);
    testRefresh(tokens.refreshToken);
    sleep(0.02);
    testLogout(tokens.accessToken);
  } else if (rand < 0.75) {
    const tokens = testLogin();
    sleep(0.05);
    testRefresh(tokens.refreshToken);
  } else {
    const tokens = testLogin();
    sleep(0.05);
    testLogout(tokens.accessToken);
  }

  sleep(0.05);
}

// ── Individual endpoint scenarios for targeted testing ───────────────────────
export const endpointScenarios = {
  login_only: {
    executor: 'constant-arrival-rate',
    duration: DURATION,
    rate: TARGET_RPS,
    preAllocatedVUs: 20,
    maxVUs: 100,
    exec: 'loginOnly',
    startTime: '0s',
    tags: { scenario: 'login_only' },
  },
  logout_only: {
    executor: 'constant-arrival-rate',
    duration: DURATION,
    rate: TARGET_RPS,
    preAllocatedVUs: 20,
    maxVUs: 100,
    exec: 'logoutOnly',
    startTime: '0s',
    tags: { scenario: 'logout_only' },
  },
  refresh_only: {
    executor: 'constant-arrival-rate',
    duration: DURATION,
    rate: TARGET_RPS,
    preAllocatedVUs: 20,
    maxVUs: 100,
    exec: 'refreshOnly',
    startTime: '0s',
    tags: { scenario: 'refresh_only' },
  },
};

export function loginOnly() {
  testLogin();
  sleep(0.1);
}

export function logoutOnly() {
  const tokens = testLogin();
  sleep(0.05);
  testLogout(tokens.accessToken);
  sleep(0.05);
}

export function refreshOnly() {
  const tokens = testLogin();
  sleep(0.05);
  testRefresh(tokens.refreshToken);
  sleep(0.05);
}

// ── Summary Hook ─────────────────────────────────────────────────────────────
export function handleSummary(data) {
  return {
    'stdout': `\n=== Darkwatch Auth Load Test Results ===\n`,
    'summary.json': JSON.stringify({
      timestamp: new Date().toISOString(),
      duration: DURATION,
      targetRPS: TARGET_RPS,
      thresholds: THRESHOLDS,
      metrics: {
        login: {
          p99: data.metrics.login_p99?.values['p(99)']?.toFixed(2) || 'N/A',
          p95: data.metrics.login_p99?.values['p(95)']?.toFixed(2) || 'N/A',
          avg: data.metrics.login_p99?.values.avg?.toFixed(2) || 'N/A',
          count: data.metrics.login_p99?.values.count || 0,
          successRate: (data.metrics.login_success?.values.rate || 0) * 100 + '%',
        },
        logout: {
          p99: data.metrics.logout_p99?.values['p(99)']?.toFixed(2) || 'N/A',
          p95: data.metrics.logout_p99?.values['p(95)']?.toFixed(2) || 'N/A',
          avg: data.metrics.logout_p99?.values.avg?.toFixed(2) || 'N/A',
          count: data.metrics.logout_p99?.values.count || 0,
          successRate: (data.metrics.logout_success?.values.rate || 0) * 100 + '%',
        },
        refresh: {
          p99: data.metrics.refresh_p99?.values['p(99)']?.toFixed(2) || 'N/A',
          p95: data.metrics.refresh_p99?.values['p(95)']?.toFixed(2) || 'N/A',
          avg: data.metrics.refresh_p99?.values.avg?.toFixed(2) || 'N/A',
          count: data.metrics.refresh_p99?.values.count || 0,
          successRate: (data.metrics.refresh_success?.values.rate || 0) * 100 + '%',
        },
      },
      passed: Object.entries(data.metrics).every(
        ([_, metric]) => metric?.thresholds?.every?.((t) => t.pass)
      ),
    }, null, 2),
  };
}