feat(vitest-plugin): add support for benchmark hooks

Author: adriencaccia

examples/with-typescript-esm/package.json

@@ -16,6 +16,6 @@
     "esbuild-register": "^3.4.2",
     "tinybench": "^2.5.0",
     "typescript": "^5.1.3",
-    "vitest": "^1.0.3"
+    "vitest": "^1.2.2"
   }
 }

packages/vitest-plugin/README.md

@@ -20,7 +20,7 @@ First, install the plugin [`@codspeed/vitest-plugin`](https://www.npmjs.com/pack
 
 > [!NOTE]
 > The CodSpeed plugin is only compatible with
-> [vitest@1.0.0](https://www.npmjs.com/package/vitest/v/1.0.0)
+> [vitest@1.2.2](https://www.npmjs.com/package/vitest/v/1.2.2)
 > and above.
 
 ```sh

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

@@ -0,0 +1,92 @@
+import {
+  afterAll,
+  afterEach,
+  beforeAll,
+  beforeEach,
+  bench,
+  describe,
+  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("run", () => {
+    beforeAll(() => {
+      cleanUpCount += 10;
+    });
+    beforeEach(() => {
+      cleanUpCount += 1;
+    });
+    afterEach(() => {
+      cleanUpCount -= 1;
+    });
+    afterAll(() => {
+      cleanUpCount -= 10;
+    });
+
+    bench("one", () => {
+      expect(cleanUpCount).toBe(11);
+    });
+    bench("two", () => {
+      expect(cleanUpCount).toBe(11);
+    });
+  });
+  bench("end", () => {
+    expect(cleanUpCount).toBe(0);
+  });
+});

packages/vitest-plugin/package.json

@@ -31,12 +31,12 @@
   },
   "peerDependencies": {
     "vite": "^4.2.0 || ^5.0.0",
-    "vitest": ">=1.0.0-beta.4 || >=1"
+    "vitest": ">=1.2.2"
   },
   "devDependencies": {
     "@total-typescript/shoehorn": "^0.1.1",
     "execa": "^8.0.1",
     "vite": "^5.0.0",
-    "vitest": "^1.0.3"
+    "vitest": "^1.2.2"
   }
 }

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

@@ -1,5 +1,5 @@
 import { fromPartial } from "@total-typescript/shoehorn";
-import { describe, expect, it, vi } from "vitest";
+import { describe, expect, it, Suite, vi } from "vitest";
 import { getBenchFn } from "vitest/suite";
 import CodSpeedRunner from "../runner";
 
@@ -25,21 +25,31 @@ vi.mock("@codspeed/core", async (importOriginal) => {
 
 console.log = vi.fn();
 
-vi.mock("vitest/suite");
+vi.mock("vitest/suite", () => ({
+  getBenchFn: vi.fn(),
+  // wrapping the value in vi.fn(...) here will not work for some reason
+  getHooks: () => ({
+    beforeAll: [],
+    beforeEach: [],
+    afterEach: [],
+    afterAll: [],
+  }),
+}));
 const mockedGetBenchFn = vi.mocked(getBenchFn);
+
 describe("CodSpeedRunner", () => {
   it("should run the bench functions only twice", async () => {
     const benchFn = vi.fn();
     mockedGetBenchFn.mockReturnValue(benchFn);
 
     const runner = new CodSpeedRunner(fromPartial({}));
-    await runner.runSuite(
-      fromPartial({
-        filepath: __filename,
-        name: "test suite",
-        tasks: [{ mode: "run", meta: { benchmark: true }, name: "test bench" }],
-      })
-    );
+    const suite = fromPartial<Suite>({
+      filepath: __filename,
+      name: "test suite",
+      tasks: [{ mode: "run", meta: { benchmark: true }, name: "test bench" }],
+    });
+    suite.tasks[0].suite = suite;
+    await runner.runSuite(suite);
 
     // setup
     expect(coreMocks.setupCore).toHaveBeenCalledTimes(1);
@@ -71,26 +81,29 @@ describe("CodSpeedRunner", () => {
     mockedGetBenchFn.mockReturnValue(benchFn);
 
     const runner = new CodSpeedRunner(fromPartial({}));
-    await runner.runSuite(
-      fromPartial({
-        filepath: __filename,
-        name: "test suite",
-        tasks: [
-          {
-            type: "suite",
-            name: "nested suite",
-            mode: "run",
-            tasks: [
-              {
-                mode: "run",
-                meta: { benchmark: true },
-                name: "test bench",
-              },
-            ],
-          },
-        ],
-      })
-    );
+    const rootSuite = fromPartial<Suite>({
+      filepath: __filename,
+      name: "test suite",
+      tasks: [
+        {
+          type: "suite",
+          name: "nested suite",
+          mode: "run",
+          tasks: [
+            {
+              mode: "run",
+              meta: { benchmark: true },
+              name: "test bench",
+            },
+          ],
+        },
+      ],
+    });
+    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);
 
     // setup
     expect(coreMocks.setupCore).toHaveBeenCalledTimes(1);

packages/vitest-plugin/src/runner.ts

@@ -8,9 +8,33 @@ import {
   teardownCore,
 } from "@codspeed/core";
 import path from "path";
-import { Benchmark, Suite } from "vitest";
+import { Benchmark, chai, Suite, Task } from "vitest";
 import { NodeBenchmarkRunner } from "vitest/runners";
-import { getBenchFn } from "vitest/suite";
+import { getBenchFn, getHooks } from "vitest/suite";
+
+type SuiteHooks = ReturnType<typeof getHooks>;
+
+function getSuiteHooks(suite: Suite, name: keyof SuiteHooks) {
+  return getHooks(suite)[name];
+}
+
+export async function callSuiteHook<T extends keyof SuiteHooks>(
+  suite: Suite,
+  currentTask: Task,
+  name: T
+): Promise<void> {
+  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"
@@ -26,39 +50,62 @@ function logCodSpeed(message: string) {
   console.log(`[CodSpeed] ${message}`);
 }
 
-async function runBenchmarkSuite(suite: Suite, parentSuiteName?: string) {
-  const benchmarkGroup: Benchmark[] = [];
-  const benchmarkSuiteGroup: Suite[] = [];
-  for (const task of suite.tasks) {
-    if (task.mode !== "run") continue;
+async function runBench(benchmark: Benchmark, currentSuiteName: string) {
+  const uri = `${currentSuiteName}::${benchmark.name}`;
+  const fn = getBenchFn(benchmark);
 
-    if (task.meta?.benchmark) benchmarkGroup.push(task as Benchmark);
-    else if (task.type === "suite") benchmarkSuiteGroup.push(task);
+  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 callSuiteHook(benchmark.suite, benchmark, "beforeEach");
+  await mongoMeasurement.start(uri);
+  await (async function __codspeed_root_frame__() {
+    Measurement.startInstrumentation();
+    // @ts-expect-error we do not need to bind the function to an instance of tinybench's Bench
+    await fn();
+    Measurement.stopInstrumentation(uri);
+  })();
+  await mongoMeasurement.stop(uri);
+  await callSuiteHook(benchmark.suite, benchmark, "afterEach");
+
+  logCodSpeed(`${uri} done`);
+}
 
+async function runBenchmarkSuite(suite: Suite, parentSuiteName?: string) {
   const currentSuiteName = parentSuiteName
     ? parentSuiteName + "::" + suite.name
     : suite.name;
 
-  for (const subSuite of benchmarkSuiteGroup) {
-    await runBenchmarkSuite(subSuite, currentSuiteName);
+  // 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");
   }
 
-  for (const benchmark of benchmarkGroup) {
-    const uri = `${currentSuiteName}::${benchmark.name}`;
-    const fn = getBenchFn(benchmark);
+  for (const task of suite.tasks) {
+    if (task.mode !== "run") continue;
 
-    await optimizeFunction(fn);
-    await mongoMeasurement.start(uri);
-    await (async function __codspeed_root_frame__() {
-      Measurement.startInstrumentation();
-      // @ts-expect-error we do not need to bind the function to an instance of tinybench's Bench
-      await fn();
-      Measurement.stopInstrumentation(uri);
-    })();
-    await mongoMeasurement.stop(uri);
-
-    logCodSpeed(`${uri} done`);
+    if (task.meta?.benchmark) {
+      await runBench(task as Benchmark, currentSuiteName);
+    } else if (task.type === "suite") {
+      await runBenchmarkSuite(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");
   }
 }