pi-swarm/manifests/loadtest/k6-tokamak-script.yaml

apiVersion: v1
kind: ConfigMap
metadata:
  name: k6-tokamak-script
  namespace: loadtest
data:
  tokamak.js: |
    // Staircase load test for the tokamak public API.
    // Setup harvests real IDs from the list endpoints, then iterations fan out
    // across the whole endpoint collection, weighted so light calls dominate
    // and the two heavy data endpoints are sampled less often.
    //
    // Env: BASE_URL, TOKAMAK_API_KEY, START_RPS, STEP_RPS, STEP_DURATION, STEPS
    import http from 'k6/http';
    import { check, fail } from 'k6';
    import exec from 'k6/execution';

    const BASE = __ENV.BASE_URL;
    const auth = { headers: { Authorization: 'Bearer ' + __ENV.TOKAMAK_API_KEY } };

    const startRps = parseInt(__ENV.START_RPS || '5');
    const stepRps = parseInt(__ENV.STEP_RPS || '5');
    const stepDurMs = ms(__ENV.STEP_DURATION || '1m');
    const steps = parseInt(__ENV.STEPS || '4');
    const rampMs = 5000;

    function ms(s) {
      const m = /^(\d+)(ms|s|m)$/.exec(String(s));
      if (!m) throw new Error('bad duration: ' + s);
      return +m[1] * ({ ms: 1, s: 1000, m: 60000 })[m[2]];
    }
    function ms2str(n) { return n + 'ms'; }

    // staircase stages + per-step windows (5s ramp + hold each)
    const stages = [];
    const windows = [];
    let clock = 0;
    for (let i = 0; i < steps; i++) {
      const rate = startRps + i * stepRps;
      stages.push({ duration: ms2str(rampMs), target: rate });
      clock += rampMs;
      const wStart = clock;
      stages.push({ duration: ms2str(stepDurMs), target: rate });
      clock += stepDurMs;
      windows.push([wStart, clock, 'step' + (i + 1) + ' @ ' + rate + ' rps']);
    }
    const totalMs = clock;
    const maxRate = startRps + (steps - 1) * stepRps;
    const vus = Math.max(10, Math.ceil(maxRate * 1.5));

    export const options = {
      discardResponseBodies: false,
      scenarios: {
        staircase: {
          executor: 'ramping-arrival-rate',
          startRate: 0,
          preAllocatedVUs: vus,
          maxVUs: vus,
          stages: stages,
        },
      },
      thresholds: { http_req_failed: ['rate<0.05'] },
    };

    // defensive id harvesting - handle data/items/results/root-array shapes
    function idsOf(json) {
      if (!json) return [];
      const arr = json.data || json.items || json.results
        || json.brands || json.cohorts || json.metrics || json.frameworks || json;
      if (!Array.isArray(arr)) return [];
      return arr.map(function (x) {
        return x && (x.id || x.brand_id || x.cohort_id || x.metric_id || x.framework_id);
      }).filter(Boolean);
    }

    export function setup() {
      const me = http.get(BASE + '/me', auth);
      if (me.status !== 200) {
        fail('auth check failed: /me returned ' + me.status + ' - check API key and environment');
      }
      const brands = idsOf(http.get(BASE + '/brands?per_page=100', auth).json());
      const cohorts = idsOf(http.get(BASE + '/cohorts?per_page=100', auth).json());
      const metrics = idsOf(http.get(BASE + '/metrics', auth).json());
      const frameworks = idsOf(http.get(BASE + '/frameworks', auth).json());
      console.log('harvested: brands=' + brands.length + ' cohorts=' + cohorts.length
        + ' metrics=' + metrics.length + ' frameworks=' + frameworks.length);
      if (!brands.length) fail('no brand ids harvested - cannot build data calls');
      return { brands: brands, cohorts: cohorts, metrics: metrics, frameworks: frameworks };
    }

    // weighted call pool: light lists and detail reads, heavy data sampled less
    function pickCall(d, i) {
      const b = d.brands[i % d.brands.length];
      const c = d.cohorts.length ? d.cohorts[i % d.cohorts.length] : null;
      const m = d.metrics.length ? d.metrics[i % d.metrics.length] : null;
      const f = d.frameworks.length ? d.frameworks[i % d.frameworks.length] : null;
      const r = Math.random();
      if (r < 0.10) return { name: 'getStatus', path: '/' };
      if (r < 0.20) return { name: 'getMe', path: '/me' };
      if (r < 0.32) return { name: 'listBrands', path: '/brands?per_page=25' };
      if (r < 0.42) return { name: 'listMetrics', path: '/metrics' };
      if (r < 0.50) return { name: 'listFrameworks', path: '/frameworks' };
      if (r < 0.62) return { name: 'getBrand', path: '/brands/' + b };
      if (r < 0.70) return c
        ? { name: 'getCohort', path: '/cohorts/' + c }
        : { name: 'listBrands', path: '/brands?per_page=25' };
      if (r < 0.76) return c
        ? { name: 'listCohortBrands', path: '/cohorts/' + c + '/brands?per_page=25' }
        : { name: 'listCohorts', path: '/cohorts?per_page=25' };
      if (r < 0.78) return f
        ? { name: 'getFramework', path: '/frameworks/' + f }
        : { name: 'listFrameworks', path: '/frameworks' };
      if (r < 0.90) return m
        ? { name: 'getMetricData', path: '/metric/data/' + m + '/' + b }
        : { name: 'listMetrics', path: '/metrics' };
      return f
        ? { name: 'getFrameworkData', path: '/framework/data/' + f + '/' + b }
        : { name: 'listFrameworks', path: '/frameworks' };
    }

    export default function (d) {
      const elapsed = exec.scenario.progress * totalMs;
      let step = 'ramp-up';
      for (const w of windows) {
        if (elapsed >= w[0] && elapsed < w[1]) { step = w[2]; break; }
      }
      const call = pickCall(d, exec.scenario.iterationInTest);
      const res = http.get(BASE + call.path, {
        headers: auth.headers,
        tags: { name: call.name, step: step },
        timeout: '30s',
      });
      check(res, { ok: function (r) { return r.status < 300; } });
    }