feat(tinybench-plugin): support afterAll, afterEach, beforeAll and be…

…foreEach
CodSpeedHQ · Jul 31, 2023 · 43d1618 · 43d1618
1 parent 06cb38d
commit 43d1618
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 0 deletions.
diff --git a/packages/tinybench-plugin/src/index.ts b/packages/tinybench-plugin/src/index.ts
@@ -47,6 +47,15 @@ export function withCodSpeed(bench: Bench): Bench {
     setupCore();
     console.log(`[CodSpeed] running with @codspeed/tinybench v${__VERSION__}`);
     for (const task of bench.tasks) {
+      // run before hooks
+      if (task.opts.beforeAll != null) {
+        await task.opts.beforeAll.call(task);
+      }
+      if (task.opts.beforeEach != null) {
+        await task.opts.beforeEach.call(task);
+      }
+
+      // run the actual benchmark, with instrumentation
       const uri = isCodSpeedBenchOptions(task.opts)
         ? task.opts.uri
         : `${rootCallingFile}::${task.name}`;
@@ -56,6 +65,16 @@ export function withCodSpeed(bench: Bench): Bench {
         await task.fn();
         Measurement.stopInstrumentation(uri);
       })();
+
+      // run after hooks
+      if (task.opts.afterEach != null) {
+        await task.opts.afterEach.call(task);
+      }
+      if (task.opts.afterAll != null) {
+        await task.opts.afterAll.call(task);
+      }
+
+      // print results
       console.log(`    ✔ Measured ${uri}`);
     }
     console.log(`[CodSpeed] Done running ${bench.tasks.length} benches.`);

diff --git a/packages/tinybench-plugin/tests/index.integ.test.ts b/packages/tinybench-plugin/tests/index.integ.test.ts
@@ -131,4 +131,34 @@ describe("Benchmark.Suite", () => {
       );
     }
   );
+
+  it("should run before and after hooks", async () => {
+    mockCore.isInstrumented.mockReturnValue(true);
+    const beforeAll = jest.fn();
+    const beforeEach = jest.fn();
+    const afterEach = jest.fn();
+    const afterAll = jest.fn();
+
+    await withCodSpeed(new Bench())
+      .add(
+        "RegExp",
+        function () {
+          /o/.test("Hello World!");
+        },
+        { afterAll, afterEach, beforeAll, beforeEach }
+      )
+      .add(
+        "RegExp2",
+        () => {
+          /o/.test("Hello World!");
+        },
+        { afterAll, afterEach, beforeAll, beforeEach }
+      )
+      .run();
+
+    expect(beforeAll).toHaveBeenCalledTimes(2);
+    expect(beforeEach).toHaveBeenCalledTimes(2);
+    expect(afterEach).toHaveBeenCalledTimes(2);
+    expect(afterAll).toHaveBeenCalledTimes(2);
+  });
 });