diff --git a/packages/vitest-plugin/src/__tests__/instrumented.test.ts b/packages/vitest-plugin/src/__tests__/instrumented.test.ts index b749b236..443d62f2 100644 --- a/packages/vitest-plugin/src/__tests__/instrumented.test.ts +++ b/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({ + 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", + ]); + }); }); diff --git a/packages/vitest-plugin/src/analysis.ts b/packages/vitest-plugin/src/analysis.ts index ff1a6a57..dc50a5e2 100644 --- a/packages/vitest-plugin/src/analysis.ts +++ b/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();