@indigo-labs/indigo-sdk/tests/setup.ts

Indigo SDK for interacting with Indigo endpoints via lucid-evolution

import { BenchmarkResult } from './utils/benchmark-utils';
import { writeFileSync, readFileSync, existsSync, unlinkSync } from 'fs';
import { join } from 'path';

// File path to persist benchmark results across workers
const BENCHMARK_RESULTS_FILE = join(
  process.cwd(),
  '.vitest-benchmark-results.json',
);

// Parse command-line arguments for --benchmark-file flag
function getBenchmarkOutputFile(): string | null {
  const args = process.argv;
  const flagIndex = args.indexOf('--benchmark-file');
  if (flagIndex !== -1 && flagIndex + 1 < args.length) {
    return args[flagIndex + 1];
  }
  return null;
}

// Helper function to read results from file
function readResultsFromFile(): Record<string, BenchmarkResult> {
  if (existsSync(BENCHMARK_RESULTS_FILE)) {
    try {
      const content = readFileSync(BENCHMARK_RESULTS_FILE, 'utf-8');
      return JSON.parse(content) as Record<string, BenchmarkResult>;
    } catch (error) {
      console.warn('Failed to read benchmark results file:', error);
      return {};
    }
  }
  return {};
}

// Helper function to write results to file
function writeResultsToFile(results: Record<string, BenchmarkResult>): void {
  try {
    writeFileSync(
      BENCHMARK_RESULTS_FILE,
      JSON.stringify(results, null, 2),
      'utf-8',
    );
  } catch (error) {
    console.warn('Failed to write benchmark results file:', error);
  }
}

// In-memory store for the current worker/test file
// Using a Proxy to automatically sync to file when properties are set
const _BENCHMARK_RESULTS: Record<string, BenchmarkResult> = {};

// Create a Proxy that automatically syncs to file on property assignment
export const BENCHMARK_RESULTS = new Proxy(_BENCHMARK_RESULTS, {
  set(target, property, value: BenchmarkResult) {
    if (Object.prototype.hasOwnProperty.call(target, property)) {
      throw new Error(
        `Benchmark with name "${property as string}" already exists.`,
      );
    }

    // Set the value in the local object
    target[property as string] = value;

    // Merge with existing results from file and write back
    const existingResults = readResultsFromFile();
    const mergedResults = { ...existingResults, ...target };
    writeResultsToFile(mergedResults);

    return true;
  },
  get(target, property) {
    // When reading, first merge with file results to get latest from other workers
    const fileResults = readResultsFromFile();
    Object.assign(target, fileResults);
    return target[property as string];
  },
});

// Vitest global setup (must be default export)
export default function setup() {
  // Clear any existing results file at the start
  if (existsSync(BENCHMARK_RESULTS_FILE)) {
    unlinkSync(BENCHMARK_RESULTS_FILE);
  }

  // Return teardown function
  return function teardown() {
    // Read all results from file (accumulated across all workers)
    const results = readResultsFromFile();

    if (Object.keys(results).length === 0) {
      return;
    }

    // Sort entries by benchmark name alphabetically
    const sortedEntries = Object.entries(results).sort(([a], [b]) =>
      a.localeCompare(b),
    );

    const decimalPlaces = 3;

    console.log('\n=== Benchmark Results ===');
    console.table(
      sortedEntries.map(([name, result]: [string, BenchmarkResult]) => ({
        'Benchmark Name': name,
        'Memory Percentage': result.memPercentage.toFixed(decimalPlaces) + '%',
        'Steps Percentage': result.stepsPercentage.toFixed(decimalPlaces) + '%',
        'Size Percentage': result.sizePercentage.toFixed(decimalPlaces) + '%',
      })),
    );

    // Write results to output file if --benchmark-file flag is present
    const outputFile = getBenchmarkOutputFile();
    if (outputFile) {
      try {
        writeFileSync(
          outputFile,
          JSON.stringify(Object.fromEntries(sortedEntries), null, 2),
          'utf-8',
        );
        console.log(`\nBenchmark results written to: ${outputFile}`);
      } catch (error) {
        console.warn(
          `Failed to write benchmark results to ${outputFile}:`,
          error,
        );
      }
    }

    // Clean up the results file
    if (existsSync(BENCHMARK_RESULTS_FILE)) {
      unlinkSync(BENCHMARK_RESULTS_FILE);
    }
  };
}