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);
const CREDENTIAL_POOL_SIZE = parseInt(__ENV.CREDENTIAL_POOL_SIZE || '100', 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',
};

// ── P1#3: Fixed credential pool (reuses pre-seeded users, not unique per call) ──
const credentialPool = Array.from({ length: CREDENTIAL_POOL_SIZE }, (_, i) => ({
  email: `${TEST_EMAIL.replace('@', `_${i}@`)}`,
  password: TEST_PASSWORD,
}));

// Fake token pool fallback — used when setup() warmup is skipped or fails
const tokenPool = Array.from({ length: CREDENTIAL_POOL_SIZE }, () => ({
  accessToken: uuidv4(),
  refreshToken: uuidv4(),
}));

// ── Setup: Seed real tokens via login warmup ──────────────────────────────────
export function setup() {
  const creds = credentialPool[0];
  const payload = JSON.stringify({ email: creds.email, password: creds.password });
  const res = http.post(`${BASE_URL}/auth/login`, payload, { headers: authHeaders });

  try {
    const json = JSON.parse(res.body);
    const accessToken = json.access_token || json.token || json.data?.access_token;
    const refreshToken = json.refresh_token || json.data?.refresh_token;

    if (accessToken && refreshToken) {
      return {
        accessToken,
        refreshToken,
        warmupSuccess: true,
      };
    }
  } catch {
    // fall through to fake tokens
  }

  console.warn(`[warmup] Login returned ${res.status} — standalone scenarios will use fake tokens (expect 401/403)`);
  return {
    accessToken: tokenPool[0].accessToken,
    refreshToken: tokenPool[0].refreshToken,
    warmupSuccess: false,
  };
}

// ── Scenario: Login (POST /auth/login) ──────────────────────────────────────
function testLogin(email, password) {
  const creds = email
    ? { email, password }
    : credentialPool[Math.floor(Math.random() * credentialPool.length)];

  const payload = JSON.stringify({
    email: creds.email,
    password: creds.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 || tokenPool[Math.floor(Math.random() * tokenPool.length)].refreshToken;

  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,
    };
  }
}

// ── P2#4: Scenario: Logout (POST /auth/logout) — refresh_token in body, Bearer in header ──
function testLogout(accessToken, refreshToken) {
  const poolEntry = tokenPool[Math.floor(Math.random() * tokenPool.length)];
  const token = accessToken || poolEntry.accessToken;
  const refreshTkn = refreshToken || poolEntry.refreshToken;

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

  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,
  });
}

// ── P1#1 + P1#2: Options with all scenarios merged (each iteration = 1 HTTP call) ──
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' },
    },
    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' },
    },
  },
  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'],
  },
};

// P1#1: Mixed workload — exactly 1 HTTP call per iteration, weighted 40/35/25
export function mixedWorkload() {
  const rand = Math.random();

  if (rand < 0.4) {
    testLogin();
  } else if (rand < 0.75) {
    testRefresh();
  } else {
    testLogout();
  }
}

// Individual endpoint scenarios — each makes exactly 1 HTTP call per iteration
export function loginOnly() {
  testLogin();
  sleep(0.1);
}

export function logoutOnly(data) {
  if (data && data.warmupSuccess) {
    testLogout(data.accessToken, data.refreshToken);
  } else {
    const poolEntry = tokenPool[Math.floor(Math.random() * tokenPool.length)];
    console.warn('[logoutOnly] Using fake token (warmup skipped or failed)');
    testLogout(poolEntry.accessToken, poolEntry.refreshToken);
  }
  sleep(0.1);
}

export function refreshOnly(data) {
  if (data && data.warmupSuccess) {
    testRefresh(data.refreshToken);
  } else {
    const poolEntry = tokenPool[Math.floor(Math.random() * tokenPool.length)];
    console.warn('[refreshOnly] Using fake token (warmup skipped or failed)');
    testRefresh(poolEntry.refreshToken);
  }
  sleep(0.1);
}

// ── Summary Hook ─────────────────────────────────────────────────────────────
export function handleSummary(data) {
  // P2#5: Only evaluate metrics that have thresholds defined
  const thresholdedMetrics = Object.entries(data.metrics).filter(
    ([_, metric]) => metric && metric.thresholds && metric.thresholds.length > 0
  );

  const passed = thresholdedMetrics.every(([_, metric]) =>
    metric.thresholds.every((t) => t.pass)
  );

  const loginP99 = data.metrics.login_p99?.values['p(99)']?.toFixed(2) || 'N/A';
  const logoutP99 = data.metrics.logout_p99?.values['p(99)']?.toFixed(2) || 'N/A';
  const refreshP99 = data.metrics.refresh_p99?.values['p(99)']?.toFixed(2) || 'N/A';

  return {
    'stdout': `\n=== Darkwatch Auth Load Test Results ===\n` +
      `Login P99:     ${loginP99}ms (threshold: ${THRESHOLDS.login}ms)\n` +
      `Logout P99:    ${logoutP99}ms (threshold: ${THRESHOLDS.logout}ms)\n` +
      `Refresh P99:   ${refreshP99}ms (threshold: ${THRESHOLDS.refresh}ms)\n` +
      `Overall:       ${passed ? 'PASS' : 'FAIL'}\n`,
  };
}