Flake Tester

flake-test.ts

#!/usr/bin/env bun

/**
 * Flake Test Runner
 * 
 * A utility for repeatedly running potentially flaky tests to measure their reliability.
 * Each run is logged to a separate file, and exit codes are tracked to calculate statistics.
 * 
 * Requirements:
 *   - Bun runtime (https://bun.sh) - a fast JavaScript runtime and toolkit
 * 
 * Usage:
 *   bun run flake-test.ts -n <count> [-c <concurrency>] -x '<shell-command>'
 * 
 * Arguments:
 *   -n <count>         Number of times to run the command (default: 10)
 *   -c <concurrency>   Number of concurrent runs (default: 1)
 *   -x <shell-command> Shell command string; supports &&, ||, |, $VAR, etc.
 * 
 * Examples:
 *   # Run a Go test 100 times with 8 concurrent executions
 *   bun run flake-test.ts -n 100 -c 8 -x 'go test -race ./pkg/...'
 * 
 *   # Pass environment variables from outside the process
 *   RUN_MCP=1 bun run flake-test.ts -n 50 -x 'echo $RUN_MCP'
 * 
 *   # Or set them inside the command
 *   bun run flake-test.ts -n 1 -x 'RUN_MCP=1 && echo $RUN_MCP'
 * 
 *   # Use pipes, variable expansion, and shell operators
 *   bun run flake-test.ts -n 20 -x 'echo $USER | grep -q root && exit 1 || exit 0'
 * 
 *   # Commands with single quotes (escape with '\'' to end quote, add escaped quote, restart quote)
 *   bun run flake-test.ts -n 1 -x 'echo "it'\''s working"'
 * 
 * Output:
 *   - Creates a temp directory for logs (path shown at start)
 *   - Shows live progress with spinner and status indicators (✓/✗)
 *   - Displays statistics at completion (success rate, failures, etc.)
 *   - Each run's output is saved to run-NNN.log in the temp directory
 */

import { $ } from "bun";
import { parseArgs } from "util";
import { mkdtemp, mkdir } from "fs/promises";
import { tmpdir } from "os";
import { join } from "path";

const SPINNER_FRAMES = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
const GREEN = "\x1b[32m";
const RED = "\x1b[31m";
const RESET = "\x1b[0m";
const SUCCESS = `${GREEN}✓${RESET}`;
const FAILURE = `${RED}✗${RESET}`;

interface TestResult {
  run: number;
  exitCode: number;
  logFile: string;
}

async function main() {
  const { values } = parseArgs({
    args: Bun.argv.slice(2),
    options: {
      n: { type: "string", short: "n", default: "10" },
      c: { type: "string", short: "c", default: "1" },
      x: { type: "string", short: "x" },
    },
    strict: false,
  });

  const shellCommand = values.x as string | undefined;
  if (!shellCommand) {
    console.error("Usage: bun run flake-test.ts -n <count> [-c <concurrency>] -x '<shell-command>'");
    process.exit(1);
  }

  const count = parseInt(values.n as string, 10);
  const concurrency = parseInt(values.c as string, 10);
  
  const tempDir = await mkdtemp(join(tmpdir(), "flake-test-"));
  const successDir = join(tempDir, "success");
  const failDir = join(tempDir, "fail");
  
  await mkdir(successDir);
  await mkdir(failDir);
  
  console.log(`\n📁 Logs directory: ${tempDir}\n`);
  console.log(`Running command ${count} times (concurrency: ${concurrency}):\n   ${shellCommand}\n`);

  const results: TestResult[] = [];
  let spinnerFrame = 0;
  let lastOutputTime = Date.now();
  let spinnerInterval: Timer | null = null;

  const updateSpinner = () => {
    const timeSinceOutput = Date.now() - lastOutputTime;
    const frame = timeSinceOutput > 2000 ? "⏸" : SPINNER_FRAMES[spinnerFrame % SPINNER_FRAMES.length];
    
    const successes = results.filter((r) => r.exitCode === 0).length;
    const failures = results.filter((r) => r.exitCode !== 0).length;
    
    // Just show progress on a single line during execution
    process.stdout.write(`\r${frame} [${results.length}/${count}] ${SUCCESS}${successes} ${FAILURE}${failures}`);
    spinnerFrame++;
  };

  spinnerInterval = setInterval(updateSpinner, 80);

  const runTest = async (i: number) => {
    const filename = `run-${i.toString().padStart(3, "0")}.log`;
    
    try {
      const proc = Bun.spawn(["/bin/sh", "-c", shellCommand], {
        stdout: "pipe",
        stderr: "pipe",
        env: { ...process.env },
      });

      // Collect output in memory
      const chunks: Uint8Array[] = [];

      const writeOutput = async (stream: ReadableStream) => {
        const reader = stream.getReader();
        while (true) {
          const { done, value } = await reader.read();
          if (done) break;
          lastOutputTime = Date.now();
          chunks.push(value);
        }
      };

      await Promise.all([
        writeOutput(proc.stdout),
        writeOutput(proc.stderr),
      ]);

      const exitCode = await proc.exited;
      
      // Write to appropriate directory based on exit code
      const targetDir = exitCode === 0 ? successDir : failDir;
      const logFile = join(targetDir, filename);
      await Bun.write(logFile, new Blob(chunks));

      results.push({ run: i, exitCode, logFile });
    } catch (error) {
      const logFile = join(failDir, filename);
      await Bun.write(logFile, "Error running test\n");
      results.push({ run: i, exitCode: 1, logFile });
    }
  };

  const workers: Promise<void>[] = [];
  let nextRun = 1;

  for (let i = 0; i < concurrency; i++) {
    workers.push((async () => {
      while (nextRun <= count) {
        const runNumber = nextRun++;
        await runTest(runNumber);
      }
    })());
  }

  await Promise.all(workers);

  if (spinnerInterval) {
    clearInterval(spinnerInterval);
  }
  
  // Clear the progress line and show full status grid
  process.stdout.write('\r\x1b[K');
  
  const terminalWidth = process.stdout.columns || 80;
  const allStatuses = results
    .sort((a, b) => a.run - b.run)
    .map((r) => r.exitCode === 0 ? SUCCESS : FAILURE);
  
  // Print status grid wrapped to terminal width
  const lines: string[] = [];
  for (let i = 0; i < allStatuses.length; i += terminalWidth) {
    lines.push(allStatuses.slice(i, i + terminalWidth).join(""));
  }
  
  console.log(lines.join('\n') + "\n");

  const successes = results.filter((r) => r.exitCode === 0).length;
  const failures = results.filter((r) => r.exitCode !== 0).length;
  const successRate = ((successes / count) * 100).toFixed(1);

  // Compress success logs
  if (successes > 0) {
    const tarFile = join(tempDir, "success.tar.gz");
    const tarProc = Bun.spawn(["tar", "-czf", tarFile, "-C", tempDir, "success"], {
      stdout: "pipe",
      stderr: "pipe",
    });
    await tarProc.exited;
    
    // Remove the success directory after compression
    await Bun.spawn(["rm", "-rf", successDir], { stdout: "pipe", stderr: "pipe" }).exited;
  }

  console.log("═".repeat(60));
  console.log(`\n📊 Results:\n`);
  console.log(`   Total runs:    ${count}`);
  console.log(`   ${SUCCESS} Successes:   ${GREEN}${successes}${RESET}`);
  console.log(`   ${FAILURE} Failures:     ${RED}${failures}${RESET}`);
  console.log(`   Success rate:  ${successRate}%`);

  console.log(`\n📁 Logs:`);
  if (successes > 0) {
    console.log(`   ${GREEN}✓${RESET} Success logs: ${join(tempDir, "success.tar.gz")}`);
  }
  if (failures > 0) {
    console.log(`   ${RED}✗${RESET} Failed logs: ${failDir}`);
  }
  
  console.log("\n" + "═".repeat(60) + "\n");

  process.exit(failures > 0 ? 1 : 0);
}

main().catch((error) => {
  console.error("Error:", error);
  process.exit(1);
});