fix(vitest-plugin): call the bench setup and teardown hooks in analysis mode

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

@@ -1,7 +1,7 @@
 import { fromPartial } from "@total-typescript/shoehorn";
 import { describe, expect, it, vi, type RunnerTestSuite } from "vitest";
 import { AnalysisRunner as CodSpeedRunner } from "../analysis";
-import { getBenchFn } from "../compat";
+import { getBenchFn, getBenchOptions } from "../compat";
 
 const coreMocks = vi.hoisted(() => {
   return {
@@ -33,9 +33,11 @@ vi.mock("../compat", async (importOriginal) => {
   return {
     ...actual,
     getBenchFn: vi.fn(),
+    getBenchOptions: vi.fn(),
   };
 });
 const mockedGetBenchFn = vi.mocked(getBenchFn);
+const mockedGetBenchOptions = vi.mocked(getBenchOptions);
 
 describe("CodSpeedRunner", () => {
   it("should run the bench function", async () => {
@@ -133,4 +135,50 @@ describe("CodSpeedRunner", () => {
     );
     expect(coreMocks.teardownCore).toHaveBeenCalledTimes(1);
   });
+  it("should call the bench setup and teardown hooks around each cycle", async () => {
+    const calls: string[] = [];
+    mockedGetBenchFn.mockReturnValue(() => {
+      calls.push("fn");
+    });
+    mockedGetBenchOptions.mockReturnValue({
+      setup: (task, mode) => {
+        calls.push(`setup:${mode}:${task.name}`);
+      },
+      teardown: (task, mode) => {
+        calls.push(`teardown:${mode}:${task.name}`);
+      },
+    });
+    coreMocks.InstrumentHooks.startBenchmark.mockImplementation(() => {
+      calls.push("startBenchmark");
+    });
+    coreMocks.InstrumentHooks.stopBenchmark.mockImplementation(() => {
+      calls.push("stopBenchmark");
+    });
+
+    const runner = new CodSpeedRunner(fromPartial({}));
+    const suite = fromPartial<RunnerTestSuite>({
+      file: { filepath: __filename },
+      name: "test suite",
+      tasks: [
+        {
+          type: "test",
+          mode: "run",
+          meta: { benchmark: true },
+          name: "test bench",
+        },
+      ],
+    });
+
+    await runner.runSuite(suite);
+
+    const hookCalls = calls.filter((call) => call !== "fn");
+    expect(hookCalls).toEqual([
+      "setup:warmup:test bench",
+      "teardown:warmup:test bench",
+      "setup:run:test bench",
+      "startBenchmark",
+      "stopBenchmark",
+      "teardown:run:test bench",
+    ]);
+  });
 });

packages/vitest-plugin/src/analysis.ts

@@ -7,13 +7,16 @@ import {
   teardownCore,
   wrapWithRootFrame,
 } from "@codspeed/core";
+import type * as tinybench from "tinybench";
 import { Benchmark, type RunnerTestSuite } from "vitest";
 import {
   callSuiteHook,
   isVitestTaskBenchmark,
   patchRootSuiteWithFullFilePath,
 } from "./common";
-import { getBenchFn, NodeBenchmarkRunner } from "./compat";
+import { getBenchFn, getBenchOptions, NodeBenchmarkRunner } from "./compat";
+
+type Tinybench = typeof tinybench;
 
 const currentFileName =
   typeof __filename === "string"
@@ -33,17 +36,27 @@ async function runAnalysisBench(
   benchmark: Benchmark,
   suite: RunnerTestSuite,
   currentSuiteName: string,
+  tinybenchModule: Tinybench,
 ) {
   const uri = `${currentSuiteName}::${benchmark.name}`;
   const fn = getBenchFn(benchmark);
 
+  // Constructing a Bench applies tinybench's no-op defaults for the setup and
+  // teardown hooks and gives them the Task they expect. The bench itself is
+  // never run: this runner drives the benchmark function directly.
+  const bench = new tinybenchModule.Bench(getBenchOptions(benchmark));
+  const task = new tinybenchModule.Task(bench, benchmark.name, fn);
+
+  await bench.setup(task, "warmup");
   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 bench.teardown(task, "warmup");
 
+  await bench.setup(task, "run");
   await callSuiteHook(suite, benchmark, "beforeEach");
   await mongoMeasurement.start(uri);
   global.gc?.();
@@ -56,12 +69,14 @@ async function runAnalysisBench(
   })();
   await mongoMeasurement.stop(uri);
   await callSuiteHook(suite, benchmark, "afterEach");
+  await bench.teardown(task, "run");
 
   logCodSpeed(`${uri} done`);
 }
 
 async function runAnalysisBenchmarkSuite(
   suite: RunnerTestSuite,
+  tinybenchModule: Tinybench,
   parentSuiteName?: string,
 ) {
   const currentSuiteName = parentSuiteName
@@ -74,9 +89,9 @@ async function runAnalysisBenchmarkSuite(
     if (task.mode !== "run") continue;
 
     if (isVitestTaskBenchmark(task)) {
-      await runAnalysisBench(task, suite, currentSuiteName);
+      await runAnalysisBench(task, suite, currentSuiteName, tinybenchModule);
     } else if (task.type === "suite") {
-      await runAnalysisBenchmarkSuite(task, currentSuiteName);
+      await runAnalysisBenchmarkSuite(task, tinybenchModule, currentSuiteName);
     }
   }
 
@@ -91,7 +106,7 @@ export class AnalysisRunner extends NodeBenchmarkRunner {
     patchRootSuiteWithFullFilePath(suite);
 
     logCodSpeed(`running suite ${suite.name}`);
-    await runAnalysisBenchmarkSuite(suite);
+    await runAnalysisBenchmarkSuite(suite, await this.importTinybench());
     logCodSpeed(`running suite ${suite.name} done`);
 
     teardownCore();