@codspeed/vitest-plugin/dist/runner.mjs

vitest plugin for CodSpeed

import { logDebug, setupCore, teardownCore, getGitDir, optimizeFunction, mongoMeasurement, Measurement } from '@codspeed/core';
import path from 'path';
import { chai } from 'vitest';
import { NodeBenchmarkRunner } from 'vitest/runners';
import { getHooks, getBenchFn } from 'vitest/suite';

function getSuiteHooks(suite, name) {
  return getHooks(suite)?.[name] ?? [];
}
async function callSuiteHook(suite, currentTask, name) {
  if (name === "beforeEach" && suite?.suite) {
    await callSuiteHook(suite.suite, currentTask, name);
  }
  const hooks = getSuiteHooks(suite, name);
  await Promise.all(hooks.map((fn) => fn()));
  if (name === "afterEach" && suite?.suite) {
    await callSuiteHook(suite.suite, currentTask, name);
  }
}
const currentFileName = typeof __filename === "string" ? __filename : new URL("runner.mjs", import.meta.url).pathname;
function logCodSpeed(message) {
  console.log(`[CodSpeed] ${message}`);
}
async function runBench(benchmark, currentSuiteName) {
  const uri = `${currentSuiteName}::${benchmark.name}`;
  const fn = getBenchFn(benchmark);
  await callSuiteHook(benchmark.suite, benchmark, "beforeEach");
  try {
    await optimizeFunction(fn);
  } catch (e) {
    if (!(e instanceof chai.AssertionError)) {
      throw e;
    }
  }
  await callSuiteHook(benchmark.suite, benchmark, "afterEach");
  await callSuiteHook(benchmark.suite, benchmark, "beforeEach");
  await mongoMeasurement.start(uri);
  global.gc?.();
  await async function __codspeed_root_frame__() {
    Measurement.startInstrumentation();
    await fn();
    Measurement.stopInstrumentation(uri);
  }();
  await mongoMeasurement.stop(uri);
  await callSuiteHook(benchmark.suite, benchmark, "afterEach");
  logCodSpeed(`${uri} done`);
}
async function runBenchmarkSuite(suite, parentSuiteName) {
  const currentSuiteName = parentSuiteName ? parentSuiteName + "::" + suite.name : suite.name;
  if (parentSuiteName !== void 0) {
    await callSuiteHook(suite, suite, "beforeAll");
  }
  for (const task of suite.tasks) {
    if (task.mode !== "run")
      continue;
    if (task.meta?.benchmark) {
      await runBench(task, currentSuiteName);
    } else if (task.type === "suite") {
      await runBenchmarkSuite(task, currentSuiteName);
    }
  }
  if (parentSuiteName !== void 0) {
    await callSuiteHook(suite, suite, "afterAll");
  }
}
function patchRootSuiteWithFullFilePath(suite) {
  if (suite.filepath === void 0) {
    throw new Error("filepath is undefined is the root suite");
  }
  const gitDir = getGitDir(suite.filepath);
  if (gitDir === void 0) {
    throw new Error("Could not find a git repository");
  }
  suite.name = path.relative(gitDir, suite.filepath);
}
class CodSpeedRunner extends NodeBenchmarkRunner {
  async runSuite(suite) {
    logDebug(`PROCESS PID: ${process.pid} in ${currentFileName}`);
    setupCore();
    patchRootSuiteWithFullFilePath(suite);
    logCodSpeed(`running suite ${suite.name}`);
    await runBenchmarkSuite(suite);
    logCodSpeed(`running suite ${suite.name} done`);
    teardownCore();
  }
}

export { callSuiteHook, CodSpeedRunner as default };