--wip-- [skip ci]

Author: adriencaccia

packages/vitest-plugin/benches/hooks.bench.ts

@@ -18,18 +18,17 @@ beforeEach(() => {
   count += 1;
 });
 
-// the count is multiplied by 2 because the bench function is called twice with codspeed (once for the optimization and once for the actual measurement)
 bench("one", () => {
-  expect(count).toBe(1 * 2);
+  expect(count).toBe(1);
 });
 
 describe("level1", () => {
   bench("two", () => {
-    expect(count).toBe(2 * 2);
+    expect(count).toBe(2);
   });
 
   bench("three", () => {
-    expect(count).toBe(3 * 2);
+    expect(count).toBe(3);
   });
 
   describe("level 2", () => {
@@ -38,12 +37,12 @@ describe("level1", () => {
     });
 
     bench("five", () => {
-      expect(count).toBe(5 * 2);
+      expect(count).toBe(5);
     });
 
     describe("level 3", () => {
       bench("seven", () => {
-        expect(count).toBe(7 * 2);
+        expect(count).toBe(7);
       });
     });
   });
@@ -54,12 +53,12 @@ describe("level1", () => {
     });
 
     bench("one", () => {
-      expect(count).toBe(1 * 2);
+      expect(count).toBe(1);
     });
   });
 
   bench("two", () => {
-    expect(count).toBe(2 * 2);
+    expect(count).toBe(2);
   });
 });
 

packages/vitest-plugin/package.json

@@ -38,6 +38,6 @@
     "@total-typescript/shoehorn": "^0.1.1",
     "execa": "^8.0.1",
     "vite": "^5.0.0",
-    "vitest": "^1.2.2"
+    "vitest": "^3.2.4"
   }
 }

packages/vitest-plugin/src/common.ts

@@ -1,17 +1,17 @@
 import { getGitDir } from "@codspeed/core";
 import path from "path";
-import { Suite, Task } from "vitest";
+import { type RunnerTask, type RunnerTestSuite } from "vitest";
 import { getHooks } from "vitest/suite";
 
 type SuiteHooks = ReturnType<typeof getHooks>;
 
-function getSuiteHooks(suite: Suite, name: keyof SuiteHooks) {
+function getSuiteHooks(suite: RunnerTestSuite, name: keyof SuiteHooks) {
   return getHooks(suite)?.[name] ?? [];
 }
 
 export async function callSuiteHook<T extends keyof SuiteHooks>(
-  suite: Suite,
-  currentTask: Task,
+  suite: RunnerTestSuite,
+  currentTask: RunnerTask,
   name: T
 ): Promise<void> {
   if (name === "beforeEach" && suite?.suite) {
@@ -20,20 +20,18 @@ export async function callSuiteHook<T extends keyof SuiteHooks>(
 
   const hooks = getSuiteHooks(suite, name);
 
+  // @ts-expect-error TODO: add support for hooks parameters
   await Promise.all(hooks.map((fn) => fn()));
 
   if (name === "afterEach" && suite?.suite) {
     await callSuiteHook(suite.suite, currentTask, name);
   }
 }
 
-export function patchRootSuiteWithFullFilePath(suite: Suite) {
-  if (suite.filepath === undefined) {
-    throw new Error("filepath is undefined is the root suite");
-  }
-  const gitDir = getGitDir(suite.filepath);
+export function patchRootSuiteWithFullFilePath(suite: RunnerTestSuite) {
+  const gitDir = getGitDir(suite.file.filepath);
   if (gitDir === undefined) {
     throw new Error("Could not find a git repository");
   }
-  suite.name = path.relative(gitDir, suite.filepath);
+  suite.name = path.relative(gitDir, suite.file.filepath);
 }

packages/vitest-plugin/src/index.ts

@@ -8,7 +8,7 @@ import {
 } from "@codspeed/core";
 import { join } from "path";
 import { Plugin } from "vite";
-import { UserConfig } from "vitest/config";
+import { type ViteUserConfig } from "vitest/config";
 
 // get this file's directory path from import.meta.url
 const __dirname = new URL(".", import.meta.url).pathname;
@@ -45,9 +45,11 @@ export default function codspeedPlugin(): Plugin {
       return true;
     },
     enforce: "post",
-    config(): UserConfig {
+    config(): ViteUserConfig {
       const runnerFile = getRunnerFile();
-      const config: UserConfig = {
+      const runnerMode = getCodspeedRunnerMode();
+
+      const config: ViteUserConfig = {
         test: {
           pool: "forks",
           poolOptions: {
@@ -56,14 +58,17 @@ export default function codspeedPlugin(): Plugin {
             },
           },
           globalSetup: [getCodSpeedFileFromName("globalSetup")],
+          ...(runnerFile && {
+            runner: runnerFile,
+          }),
+          ...(runnerMode === "walltime" && {
+            benchmark: {
+              includeSamples: true,
+            },
+          }),
         },
       };
 
-      // Only set custom runner when CODSPEED_ENV is set
-      if (runnerFile) {
-        config.test!.runner = runnerFile;
-      }
-
       return config;
     },
   };

packages/vitest-plugin/src/instrumented.ts

@@ -6,7 +6,7 @@ import {
   setupCore,
   teardownCore,
 } from "@codspeed/core";
-import { Benchmark, chai, Suite } from "vitest";
+import { Benchmark, type RunnerTestSuite } from "vitest";
 import { NodeBenchmarkRunner } from "vitest/runners";
 import { getBenchFn } from "vitest/suite";
 import { callSuiteHook, patchRootSuiteWithFullFilePath } from "./common";
@@ -32,20 +32,19 @@ async function runInstrumentedBench(
   const uri = `${currentSuiteName}::${benchmark.name}`;
   const fn = getBenchFn(benchmark);
 
-  await callSuiteHook(benchmark.suite, benchmark, "beforeEach");
-  try {
-    await optimizeFunction(fn);
-  } catch (e) {
-    // if the error is not an assertion error, we want to fail the run
-    // we allow assertion errors because we want to be able to use `expect` in the benchmark to allow for better authoring
-    // assertions are allowed to fail in the optimization phase since it might be linked to stateful code
-    if (!(e instanceof chai.AssertionError)) {
-      throw e;
+  await optimizeFunction(async () => {
+    if (benchmark.suite) {
+      await callSuiteHook(benchmark.suite, benchmark, "beforeEach");
     }
-  }
-  await callSuiteHook(benchmark.suite, benchmark, "afterEach");
+    await fn();
+    if (benchmark.suite) {
+      await callSuiteHook(benchmark.suite, benchmark, "afterEach");
+    }
+  });
 
-  await callSuiteHook(benchmark.suite, benchmark, "beforeEach");
+  if (benchmark.suite) {
+    await callSuiteHook(benchmark.suite, benchmark, "beforeEach");
+  }
   await mongoMeasurement.start(uri);
   global.gc?.();
   await (async function __codspeed_root_frame__() {
@@ -55,24 +54,22 @@ async function runInstrumentedBench(
     Measurement.stopInstrumentation(uri);
   })();
   await mongoMeasurement.stop(uri);
-  await callSuiteHook(benchmark.suite, benchmark, "afterEach");
+  if (benchmark.suite) {
+    await callSuiteHook(benchmark.suite, benchmark, "afterEach");
+  }
 
   logCodSpeed(`${uri} done`);
 }
 
 async function runInstrumentedBenchmarkSuite(
-  suite: Suite,
+  suite: RunnerTestSuite,
   parentSuiteName?: string
 ) {
   const currentSuiteName = parentSuiteName
     ? parentSuiteName + "::" + suite.name
     : suite.name;
 
-  // do not call `beforeAll` if we are in the root suite, since it is already called by vitest
-  // see https://github.com/vitest-dev/vitest/blob/1fee63f2598edc228017f18eca325f85ee54aee0/packages/runner/src/run.ts#L293
-  if (parentSuiteName !== undefined) {
-    await callSuiteHook(suite, suite, "beforeAll");
-  }
+  await callSuiteHook(suite, suite, "beforeAll");
 
   for (const task of suite.tasks) {
     if (task.mode !== "run") continue;
@@ -84,15 +81,11 @@ async function runInstrumentedBenchmarkSuite(
     }
   }
 
-  // do not call `afterAll` if we are in the root suite, since it is already called by vitest
-  // see https://github.com/vitest-dev/vitest/blob/1fee63f2598edc228017f18eca325f85ee54aee0/packages/runner/src/run.ts#L324
-  if (parentSuiteName !== undefined) {
-    await callSuiteHook(suite, suite, "afterAll");
-  }
+  await callSuiteHook(suite, suite, "afterAll");
 }
 
 export class InstrumentedRunner extends NodeBenchmarkRunner {
-  async runSuite(suite: Suite): Promise<void> {
+  async runSuite(suite: RunnerTestSuite): Promise<void> {
     logDebug(`PROCESS PID: ${process.pid} in ${currentFileName}`);
     setupCore();
 

packages/vitest-plugin/src/walltime.ts

@@ -6,27 +6,28 @@ import {
   type Benchmark,
   type BenchmarkStats,
 } from "@codspeed/core";
-import { BenchTaskResult, Suite, type Custom } from "vitest";
+import {
+  type Benchmark as VitestBenchmark,
+  type RunnerTask,
+  type RunnerTaskResult,
+  type RunnerTestSuite,
+} from "vitest";
 import { NodeBenchmarkRunner } from "vitest/runners";
 import { getBenchOptions } from "vitest/suite";
 import { patchRootSuiteWithFullFilePath } from "./common";
 
 declare const __VERSION__: string;
 
-// Vitest's task result structure extends Tinybench with additional properties
-interface VitestTaskResult {
-  state: "pass" | "fail" | "skip" | "todo";
-  startTime: number;
-  duration?: number;
-  benchmark?: BenchTaskResult;
+function isVitestTaskBenchmark(task: RunnerTask): task is VitestBenchmark {
+  return task.type === "test" && task.meta.benchmark === true;
 }
 
 /**
  * WalltimeRunner uses Vitest's default benchmark execution
  * and extracts results from the suite after completion
  */
 export class WalltimeRunner extends NodeBenchmarkRunner {
-  async runSuite(suite: Suite): Promise<void> {
+  async runSuite(suite: RunnerTestSuite): Promise<void> {
     patchRootSuiteWithFullFilePath(suite);
 
     console.log(
@@ -52,7 +53,7 @@ export class WalltimeRunner extends NodeBenchmarkRunner {
   }
 
   private async extractBenchmarkResults(
-    suite: Suite,
+    suite: RunnerTestSuite,
     parentPath = ""
   ): Promise<Benchmark[]> {
     const benchmarks: Benchmark[] = [];
@@ -61,11 +62,7 @@ export class WalltimeRunner extends NodeBenchmarkRunner {
       : suite.name;
 
     for (const task of suite.tasks) {
-      if (
-        task.meta.benchmark &&
-        task.type === "custom" &&
-        task.result?.state === "pass"
-      ) {
+      if (isVitestTaskBenchmark(task) && task.result?.state === "pass") {
         const benchmark = await this.processBenchmarkTask(task, currentPath);
         if (benchmark) {
           benchmarks.push(benchmark);
@@ -83,7 +80,7 @@ export class WalltimeRunner extends NodeBenchmarkRunner {
   }
 
   private async processBenchmarkTask(
-    task: Custom,
+    task: VitestBenchmark,
     suitePath: string
   ): Promise<Benchmark | null> {
     const uri = `${suitePath}::${task.name}`;
@@ -99,7 +96,7 @@ export class WalltimeRunner extends NodeBenchmarkRunner {
       const benchOptions = getBenchOptions(task);
 
       const stats = this.convertVitestResultToBenchmarkStats(
-        result as VitestTaskResult,
+        result,
         benchOptions
       );
 
@@ -135,7 +132,7 @@ export class WalltimeRunner extends NodeBenchmarkRunner {
   }
 
   private convertVitestResultToBenchmarkStats(
-    result: VitestTaskResult,
+    result: RunnerTaskResult,
     benchOptions: {
       time?: number;
       warmupTime?: number;