From d08b6e8401e159cd7cc7ceae391b8e6e29633d34 Mon Sep 17 00:00:00 2001
From: Nazar Hussain <nazarhussain@gmail.com>
Date: Thu, 12 Oct 2023 17:25:39 +0200
Subject: [PATCH] Add custom matcher support in the vitest

---
 scripts/vitest/customMatchers.ts | 55 ++++++++++++++++++++++++++++++++
 types/vitest/index.d.ts          | 35 ++++++++++++++++++++
 vitest.base.config.ts            | 29 +++++++++++++++++
 3 files changed, 119 insertions(+)
 create mode 100644 scripts/vitest/customMatchers.ts
 create mode 100644 types/vitest/index.d.ts
 create mode 100644 vitest.base.config.ts

diff --git a/scripts/vitest/customMatchers.ts b/scripts/vitest/customMatchers.ts
new file mode 100644
index 000000000000..04b665bf3242
--- /dev/null
+++ b/scripts/vitest/customMatchers.ts
@@ -0,0 +1,55 @@
+// eslint-disable-next-line import/no-extraneous-dependencies
+import {expect} from "vitest";
+
+expect.extend({
+  toBeValidEpochCommittee: (
+    committee: {index: number; slot: number; validators: unknown[]},
+    {
+      committeeCount,
+      validatorsPerCommittee,
+      slotsPerEpoch,
+    }: {committeeCount: number; validatorsPerCommittee: number; slotsPerEpoch: number}
+  ) => {
+    if (committee.index < 0 || committee.index > committeeCount - 1) {
+      return {
+        message: () =>
+          `Committee index out of range. Expected between 0-${committeeCount - 1}, but got ${committee.index}`,
+        pass: false,
+      };
+    }
+
+    if (committee.slot < 0 || committee.slot > slotsPerEpoch - 1) {
+      return {
+        message: () =>
+          `Committee slot out of range. Expected between 0-${slotsPerEpoch - 1}, but got ${committee.slot}`,
+        pass: false,
+      };
+    }
+
+    if (committee.validators.length !== validatorsPerCommittee) {
+      return {
+        message: () =>
+          `Incorrect number of validators in committee. Expected ${validatorsPerCommittee}, but got ${committee.validators.length}`,
+        pass: false,
+      };
+    }
+
+    return {
+      message: () => "Committee is valid",
+      pass: true,
+    };
+  },
+  toBeWithMessage: (received: unknown, expected: unknown, message: string) => {
+    if (received === expected) {
+      return {
+        message: () => "Expected value is truthy",
+        pass: true,
+      };
+    }
+
+    return {
+      pass: false,
+      message: () => message,
+    };
+  },
+});
diff --git a/types/vitest/index.d.ts b/types/vitest/index.d.ts
new file mode 100644
index 000000000000..38ccf5252d52
--- /dev/null
+++ b/types/vitest/index.d.ts
@@ -0,0 +1,35 @@
+// eslint-disable-next-line import/no-extraneous-dependencies, @typescript-eslint/no-unused-vars
+import * as vitest from "vitest";
+
+interface CustomMatchers<R = unknown> {
+  toBeValidEpochCommittee(opts: {committeeCount: number; validatorsPerCommittee: number; slotsPerEpoch: number}): R;
+  /**
+   * @deprecated
+   * We highly recommend to not use this matcher instead use detail test case
+   * where you don't need message to explain assertion
+   *
+   * @example
+   * ```ts
+   * it("should work as expected", () => {
+   *   const a = 1;
+   *   const b = 2;
+   *   expect(a).toBeWithMessage(b, "a must be equal to b");
+   * });
+   * ```
+   * can be written as:
+   * ```ts
+   * it("a should always same as b", () => {
+   *   const a = 1;
+   *   const b = 2;
+   *   expect(a).toBe(b);
+   * });
+   * ```
+   * */
+  toBeWithMessage(expected: unknown, message: string): R;
+}
+
+declare module "vitest" {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  interface Assertion<T = any> extends CustomMatchers<T> {}
+  interface AsymmetricMatchersContaining extends CustomMatchers {}
+}
diff --git a/vitest.base.config.ts b/vitest.base.config.ts
new file mode 100644
index 000000000000..34c0d56e40d5
--- /dev/null
+++ b/vitest.base.config.ts
@@ -0,0 +1,29 @@
+import path from "node:path";
+import {defineConfig} from "vitest/config";
+const __dirname = new URL(".", import.meta.url).pathname;
+
+export default defineConfig({
+  test: {
+    setupFiles: [path.join(__dirname, "./scripts/vitest/customMatchers.ts")],
+    reporters: ["default", "hanging-process"],
+    coverage: {
+      clean: true,
+      all: false,
+      extension: [".ts"],
+      provider: "v8",
+      reporter: [["lcovonly", {file: "lcov.info"}], ["text"]],
+      reportsDirectory: "./coverage",
+      exclude: [
+        "**/*.d.ts",
+        "**/*.js",
+        "**/lib/**",
+        "**/coverage/**",
+        "**/scripts/**",
+        "**/test/**",
+        "**/types/**",
+        "**/bin/**",
+        "**/node_modules/**",
+      ],
+    },
+  },
+});