ShieldAI/load-tests/voiceprint/voiceprint.js

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

// ── Configuration ────────────────────────────────────────────────────────────
const BASE_URL = __ENV.VOICEPRINT_BASE_URL || 'http://localhost:3000';
const API_TOKEN = __ENV.API_TOKEN || 'test-token';
const DURATION = __ENV.DURATION || '300s'; // 5 minutes
const TARGET_RPS = parseInt(__ENV.TARGET_RPS || '500', 10);

// P99 latency thresholds (ms)
const THRESHOLDS = {
  enrollment: parseInt(__ENV.ENROLLMENT_P99_MS || '500', 10),
  verification: parseInt(__ENV.VERIFICATION_P99_MS || '250', 10),
  modelRetrieval: parseInt(__ENV.MODEL_RETRIEVAL_P99_MS || '100', 10),
};

// ── Custom Metrics ───────────────────────────────────────────────────────────
const enrollmentLatency = new Trend('enrollment_p99');
const verificationLatency = new Trend('verification_p99');
const modelRetrievalLatency = new Trend('model_retrieval_p99');

const enrollmentSuccess = new Rate('enrollment_success');
const verificationSuccess = new Rate('verification_success');
const modelRetrievalSuccess = new Rate('model_retrieval_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);
  });
}

// Generate a realistic audio payload (base64-encoded WAV-like buffer)
// ~3 seconds of 16kHz mono 16-bit audio = ~96KB
function generateAudioPayload() {
  const size = 96000;
  const audio = new Array(size);
  for (let i = 0; i < size; i++) {
    audio[i] = Math.floor(Math.random() * 256);
  }
  return btoa(String.fromCharCode(...audio.slice(0, 2048)));
}

const headers = {
  'Content-Type': 'application/json',
  'Authorization': `Bearer ${API_TOKEN}`,
};

// ── Scenario: Enrollment (POST /voiceprint/enroll) ──────────────────────────
function testEnrollment() {
  const payload = JSON.stringify({
    name: `voice_profile_${uuidv4()}`,
    audio: generateAudioPayload(),
  });

  const res = http.post(`${BASE_URL}/voiceprint/enroll`, payload, { headers });
  const duration = res.timings.duration;
  enrollmentLatency.add(duration);

  const success = res.status === 201;
  enrollmentSuccess.add(success);

  check(res, {
    'enrollment: status 201': (r) => r.status === 201,
    'enrollment: has enrollment.id': (r) => {
      try {
        const json = JSON.parse(r.body);
        return !!json.enrollment && !!json.enrollment.id;
      } catch {
        return false;
      }
    },
    `enrollment: P99 < ${THRESHOLDS.enrollment}ms`: (r) => duration < THRESHOLDS.enrollment,
  });

  return res.json()?.enrollment?.id || uuidv4();
}

// ── Scenario: Verification (POST /voiceprint/analyze) ───────────────────────
function testVerification() {
  const payload = JSON.stringify({
    audio: generateAudioPayload(),
  });

  const res = http.post(`${BASE_URL}/voiceprint/analyze`, payload, { headers });
  const duration = res.timings.duration;
  verificationLatency.add(duration);

  const success = res.status === 201;
  verificationSuccess.add(success);

  check(res, {
    'verification: status 201': (r) => r.status === 201,
    'verification: has analysis.id': (r) => {
      try {
        const json = JSON.parse(r.body);
        return !!json.analysis && !!json.analysis.id;
      } catch {
        return false;
      }
    },
    `verification: P99 < ${THRESHOLDS.verification}ms`: (r) => duration < THRESHOLDS.verification,
  });

  return res.json()?.analysis?.id || uuidv4();
}

// ── Scenario: Model Retrieval (GET /voiceprint/results/:id) ─────────────────
function testModelRetrieval(modelId) {
  const id = modelId || uuidv4();
  const res = http.get(`${BASE_URL}/voiceprint/results/${id}`, { headers });
  const duration = res.timings.duration;
  modelRetrievalLatency.add(duration);

  // 200 = found, 404 = not found (both valid for load testing)
  const success = res.status === 200 || res.status === 404;
  modelRetrievalSuccess.add(success);

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

// ── 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: {
    `enrollment_p99`: [`p(99)<${THRESHOLDS.enrollment}`],
    `verification_p99`: [`p(99)<${THRESHOLDS.verification}`],
    `model_retrieval_p99`: [`p(99)<${THRESHOLDS.modelRetrieval}`],
    `enrollment_success`: ['rate>0.95'],
    `verification_success`: ['rate>0.95'],
    `model_retrieval_success`: ['rate>0.95'],
    http_req_duration: [`p(95)<400`, `p(99)<500`],
    http_req_failed: ['rate<0.05'],
  },
};

// Mixed workload: 30% enrollment, 45% verification, 25% model retrieval
export function mixedWorkload() {
  const rand = Math.random();

  if (rand < 0.3) {
    const modelId = testEnrollment();
    sleep(0.1);
    testModelRetrieval(modelId);
  } else if (rand < 0.75) {
    const modelId = testVerification();
    sleep(0.05);
    testModelRetrieval(modelId);
  } else {
    testModelRetrieval();
  }

  sleep(0.05);
}

// ── Individual endpoint scenarios for targeted testing ───────────────────────
export const endpointScenarios = {
  enrollment_only: {
    executor: 'constant-arrival-rate',
    duration: DURATION,
    rate: TARGET_RPS,
    preAllocatedVUs: 20,
    maxVUs: 100,
    exec: 'enrollmentOnly',
    startTime: '0s',
    tags: { scenario: 'enrollment_only' },
  },
  verification_only: {
    executor: 'constant-arrival-rate',
    duration: DURATION,
    rate: TARGET_RPS,
    preAllocatedVUs: 20,
    maxVUs: 100,
    exec: 'verificationOnly',
    startTime: '0s',
    tags: { scenario: 'verification_only' },
  },
  model_retrieval_only: {
    executor: 'constant-arrival-rate',
    duration: DURATION,
    rate: TARGET_RPS,
    preAllocatedVUs: 20,
    maxVUs: 100,
    exec: 'modelRetrievalOnly',
    startTime: '0s',
    tags: { scenario: 'model_retrieval_only' },
  },
};

export function enrollmentOnly() {
  testEnrollment();
  sleep(0.1);
}

export function verificationOnly() {
  testVerification();
  sleep(0.05);
}

export function modelRetrievalOnly() {
  testModelRetrieval();
  sleep(0.02);
}

// ── Summary Hook ─────────────────────────────────────────────────────────────
export function handleSummary(data) {
  return {
    'stdout': `\n=== Voiceprint Load Test Results ===\n`,
    'summary.json': JSON.stringify({
      timestamp: new Date().toISOString(),
      duration: DURATION,
      targetRPS: TARGET_RPS,
      thresholds: THRESHOLDS,
      metrics: {
        enrollment: {
          p99: data.metrics.enrollment_p99?.values['p(99)']?.toFixed(2) || 'N/A',
          p95: data.metrics.enrollment_p99?.values['p(95)']?.toFixed(2) || 'N/A',
          avg: data.metrics.enrollment_p99?.values.avg?.toFixed(2) || 'N/A',
          count: data.metrics.enrollment_p99?.values.count || 0,
          successRate: (data.metrics.enrollment_success?.values.rate || 0) * 100 + '%',
        },
        verification: {
          p99: data.metrics.verification_p99?.values['p(99)']?.toFixed(2) || 'N/A',
          p95: data.metrics.verification_p99?.values['p(95)']?.toFixed(2) || 'N/A',
          avg: data.metrics.verification_p99?.values.avg?.toFixed(2) || 'N/A',
          count: data.metrics.verification_p99?.values.count || 0,
          successRate: (data.metrics.verification_success?.values.rate || 0) * 100 + '%',
        },
        modelRetrieval: {
          p99: data.metrics.model_retrieval_p99?.values['p(99)']?.toFixed(2) || 'N/A',
          p95: data.metrics.model_retrieval_p99?.values['p(95)']?.toFixed(2) || 'N/A',
          avg: data.metrics.model_retrieval_p99?.values.avg?.toFixed(2) || 'N/A',
          count: data.metrics.model_retrieval_p99?.values.count || 0,
          successRate: (data.metrics.model_retrieval_success?.values.rate || 0) * 100 + '%',
        },
      },
      passed: Object.entries(data.metrics).every(
        ([_, metric]) => metric?.thresholds?.every?.((t) => t.pass)
      ),
    }, null, 2),
  };
}