feat!(vitest-plugin): support vitest v3

Author: adriencaccia

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

@@ -8,85 +8,30 @@ import {
   expect,
 } from "vitest";
 
-let count = -1;
-
-beforeAll(() => {
-  count += 1;
-});
-
-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);
-});
-
-describe("level1", () => {
-  bench("two", () => {
-    expect(count).toBe(2 * 2);
-  });
-
-  bench("three", () => {
-    expect(count).toBe(3 * 2);
-  });
-
-  describe("level 2", () => {
-    beforeEach(() => {
-      count += 1;
-    });
-
-    bench("five", () => {
-      expect(count).toBe(5 * 2);
-    });
-
-    describe("level 3", () => {
-      bench("seven", () => {
-        expect(count).toBe(7 * 2);
-      });
-    });
-  });
-
-  describe("level 2 bench nested beforeAll", () => {
-    beforeAll(() => {
-      count = 0;
-    });
-
-    bench("one", () => {
-      expect(count).toBe(1 * 2);
-    });
-  });
-
-  bench("two", () => {
-    expect(count).toBe(2 * 2);
-  });
-});
-
-describe("hooks cleanup", () => {
-  let cleanUpCount = 0;
+describe("hooks", () => {
+  let count = 0;
   describe("run", () => {
     beforeAll(() => {
-      cleanUpCount += 10;
+      count += 10;
     });
     beforeEach(() => {
-      cleanUpCount += 1;
+      count += 1;
     });
     afterEach(() => {
-      cleanUpCount -= 1;
+      count -= 1;
     });
     afterAll(() => {
-      cleanUpCount -= 10;
+      count -= 10;
     });
 
     bench("one", () => {
-      expect(cleanUpCount).toBe(11);
+      expect(count).toBe(11);
     });
     bench("two", () => {
-      expect(cleanUpCount).toBe(11);
+      expect(count).toBe(11);
     });
   });
   bench("end", () => {
-    expect(cleanUpCount).toBe(0);
+    expect(count).toBe(0);
   });
 });

packages/vitest-plugin/package.json

@@ -31,13 +31,13 @@
     "@codspeed/core": "workspace:^4.0.1"
   },
   "peerDependencies": {
-    "vite": "^4.2.0 || ^5.0.0 || ^6.0.0",
-    "vitest": ">=1.2.2"
+    "vite": "^4.2.0 || ^5.0.0 || ^6.0.0 || ^7.0.0",
+    "vitest": ">=3.2"
   },
   "devDependencies": {
     "@total-typescript/shoehorn": "^0.1.1",
     "execa": "^8.0.1",
-    "vite": "^5.0.0",
-    "vitest": "^1.2.2"
+    "vite": "^7.0.6",
+    "vitest": "^3.2.4"
   }
 }

packages/vitest-plugin/src/__tests__/index.test.ts

@@ -86,6 +86,7 @@ describe("codSpeedPlugin", () => {
     if (typeof config !== "function")
       throw new Error("config is not a function");
 
+    // @ts-expect-error we cannot type the this of the config function
     expect(config({}, fromPartial({}))).toStrictEqual({
       test: {
         globalSetup: [

packages/vitest-plugin/src/__tests__/instrumented.test.ts

@@ -1,5 +1,5 @@
 import { fromPartial } from "@total-typescript/shoehorn";
-import { describe, expect, it, Suite, vi } from "vitest";
+import { describe, expect, it, vi, type RunnerTestSuite } from "vitest";
 import { getBenchFn } from "vitest/suite";
 import { InstrumentedRunner as CodSpeedRunner } from "../instrumented";
 
@@ -37,17 +37,23 @@ vi.mock("vitest/suite", async (importOriginal) => {
 const mockedGetBenchFn = vi.mocked(getBenchFn);
 
 describe("CodSpeedRunner", () => {
-  it("should run the bench functions only twice", async () => {
+  it("should run the bench function", async () => {
     const benchFn = vi.fn();
     mockedGetBenchFn.mockReturnValue(benchFn);
 
     const runner = new CodSpeedRunner(fromPartial({}));
-    const suite = fromPartial<Suite>({
-      filepath: __filename,
+    const suite = fromPartial<RunnerTestSuite>({
+      file: { filepath: __filename },
       name: "test suite",
-      tasks: [{ mode: "run", meta: { benchmark: true }, name: "test bench" }],
+      tasks: [
+        {
+          type: "test",
+          mode: "run",
+          meta: { benchmark: true },
+          name: "test bench",
+        },
+      ],
     });
-    suite.tasks[0].suite = suite;
     await runner.runSuite(suite);
 
     // setup
@@ -80,8 +86,8 @@ describe("CodSpeedRunner", () => {
     mockedGetBenchFn.mockReturnValue(benchFn);
 
     const runner = new CodSpeedRunner(fromPartial({}));
-    const rootSuite = fromPartial<Suite>({
-      filepath: __filename,
+    const rootSuite = fromPartial<RunnerTestSuite>({
+      file: { filepath: __filename },
       name: "test suite",
       tasks: [
         {
@@ -90,6 +96,7 @@ describe("CodSpeedRunner", () => {
           mode: "run",
           tasks: [
             {
+              type: "test",
               mode: "run",
               meta: { benchmark: true },
               name: "test bench",
@@ -98,9 +105,6 @@ describe("CodSpeedRunner", () => {
         },
       ],
     });
-    rootSuite.tasks[0].suite = rootSuite;
-    // @ts-expect-error type is not narrow enough, but it is fine
-    rootSuite.tasks[0].tasks[0].suite = rootSuite.tasks[0];
 
     await runner.runSuite(rootSuite);
 

packages/vitest-plugin/src/common.ts

@@ -1,17 +1,16 @@
 import { getGitDir } from "@codspeed/core";
 import path from "path";
-import { Suite, Task } from "vitest";
+import { Benchmark, 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 +19,22 @@ 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);
+}
+
+export function isVitestTaskBenchmark(task: RunnerTask): task is Benchmark {
+  return task.type === "test" && task.meta.benchmark === true;
 }

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,10 +6,14 @@ 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";
+import {
+  callSuiteHook,
+  isVitestTaskBenchmark,
+  patchRootSuiteWithFullFilePath,
+} from "./common";
 
 const currentFileName =
   typeof __filename === "string"
@@ -27,25 +31,20 @@ function logCodSpeed(message: string) {
 
 async function runInstrumentedBench(
   benchmark: Benchmark,
+  suite: RunnerTestSuite,
   currentSuiteName: string
 ) {
   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 callSuiteHook(benchmark.suite, benchmark, "afterEach");
+  await optimizeFunction(async () => {
+    await callSuiteHook(suite, benchmark, "beforeEach");
+    // @ts-expect-error we do not need to bind the function to an instance of tinybench's Bench
+    await fn();
+    await callSuiteHook(suite, benchmark, "afterEach");
+  });
 
-  await callSuiteHook(benchmark.suite, benchmark, "beforeEach");
+  await callSuiteHook(suite, benchmark, "beforeEach");
   await mongoMeasurement.start(uri);
   global.gc?.();
   await (async function __codspeed_root_frame__() {
@@ -55,44 +54,36 @@ async function runInstrumentedBench(
     Measurement.stopInstrumentation(uri);
   })();
   await mongoMeasurement.stop(uri);
-  await callSuiteHook(benchmark.suite, benchmark, "afterEach");
+  await callSuiteHook(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;
 
-    if (task.meta?.benchmark) {
-      await runInstrumentedBench(task as Benchmark, currentSuiteName);
+    if (isVitestTaskBenchmark(task)) {
+      await runInstrumentedBench(task, suite, currentSuiteName);
     } else if (task.type === "suite") {
       await runInstrumentedBenchmarkSuite(task, currentSuiteName);
     }
   }
 
-  // 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();