From f777824a5e81e9bec198b6b83ba79adc521aef58 Mon Sep 17 00:00:00 2001
From: Nazar Hussain <nazarhussain@gmail.com>
Date: Fri, 1 Nov 2024 17:13:09 +0100
Subject: [PATCH] Add unit tests for hierarichal layers tests

---
 .../utils/hierarchicalLayers.ts               |   3 +-
 .../utils/hierarchicalLayers.test.ts          | 113 ++++++++++++++++++
 2 files changed, 115 insertions(+), 1 deletion(-)
 create mode 100644 packages/beacon-node/test/unit/chain/historialState/utils/hierarchicalLayers.test.ts

diff --git a/packages/beacon-node/src/chain/historicalState/utils/hierarchicalLayers.ts b/packages/beacon-node/src/chain/historicalState/utils/hierarchicalLayers.ts
index 0c50c3de561..7a46b66ff81 100644
--- a/packages/beacon-node/src/chain/historicalState/utils/hierarchicalLayers.ts
+++ b/packages/beacon-node/src/chain/historicalState/utils/hierarchicalLayers.ts
@@ -93,10 +93,11 @@ export class HierarchicalLayers {
         path.push(newSlot);
       }
     }
+
     const diffSlots = [...new Set(path)];
     const snapshotSlot = diffSlots.shift();
 
-    if (!snapshotSlot) {
+    if (snapshotSlot == null) {
       throw new Error(`Can not find snapshot layer for slot=${slot}`);
     }
 
diff --git a/packages/beacon-node/test/unit/chain/historialState/utils/hierarchicalLayers.test.ts b/packages/beacon-node/test/unit/chain/historialState/utils/hierarchicalLayers.test.ts
new file mode 100644
index 00000000000..c2a326ab097
--- /dev/null
+++ b/packages/beacon-node/test/unit/chain/historialState/utils/hierarchicalLayers.test.ts
@@ -0,0 +1,113 @@
+import {describe, it, expect, beforeEach} from "vitest";
+import {HierarchicalLayers, Layers} from "../../../../../src/chain/historicalState/utils/hierarchicalLayers.js";
+import {computeStartSlotAtEpoch} from "@lodestar/state-transition";
+
+describe("HierarchicalLayers", () => {
+  const layersEpochs = "1,3,5,7";
+  let hierarchicalLayers: HierarchicalLayers;
+
+  beforeEach(() => {
+    hierarchicalLayers = HierarchicalLayers.fromString(layersEpochs);
+  });
+
+  describe("toString", () => {
+    it("should be same as initialized string", () => {
+      expect(hierarchicalLayers.toString()).toEqual("1,3,5,7");
+    });
+  });
+
+  describe("totalLayers", () => {
+    it("should return valid number of layers", () => {
+      expect(hierarchicalLayers.totalLayers).toEqual(4);
+    });
+  });
+
+  describe("getArchiveLayers", () => {
+    // Considering that there are 8 slots per epoch
+    const testCases: {title: string; slot: number; output: Layers}[] = [
+      {title: "genesis slot", slot: 0, output: {snapshotSlot: 0, diffSlots: []}},
+      {title: "slot after genesis slot", slot: 5, output: {snapshotSlot: 0, diffSlots: []}},
+      {
+        title: "slot before the epoch 1",
+        slot: computeStartSlotAtEpoch(1) - 1,
+        output: {snapshotSlot: 0, diffSlots: []},
+      },
+      {
+        title: "slot at epoch 1",
+        slot: computeStartSlotAtEpoch(1),
+        output: {snapshotSlot: 0, diffSlots: [computeStartSlotAtEpoch(1)]},
+      },
+      {
+        title: "slot after epoch 1",
+        slot: computeStartSlotAtEpoch(1) + 1,
+        output: {snapshotSlot: 0, diffSlots: [computeStartSlotAtEpoch(1)]},
+      },
+      {
+        title: "slot before epoch 2",
+        slot: computeStartSlotAtEpoch(2) - 1,
+        output: {snapshotSlot: 0, diffSlots: [computeStartSlotAtEpoch(1)]},
+      },
+      {
+        title: "slot at epoch 2",
+        slot: computeStartSlotAtEpoch(2),
+        output: {snapshotSlot: 0, diffSlots: [computeStartSlotAtEpoch(2)]},
+      },
+      {
+        title: "slot after epoch 2",
+        slot: computeStartSlotAtEpoch(2) + 1,
+        output: {snapshotSlot: 0, diffSlots: [computeStartSlotAtEpoch(2)]},
+      },
+      {
+        title: "slot before epoch 3",
+        slot: computeStartSlotAtEpoch(3) - 1,
+        output: {snapshotSlot: 0, diffSlots: [computeStartSlotAtEpoch(2)]},
+      },
+      {
+        title: "slot at epoch 3",
+        slot: computeStartSlotAtEpoch(3),
+        output: {snapshotSlot: 0, diffSlots: [computeStartSlotAtEpoch(3)]},
+      },
+      {
+        title: "slot after epoch 3",
+        slot: computeStartSlotAtEpoch(3) + 1,
+        output: {snapshotSlot: 0, diffSlots: [computeStartSlotAtEpoch(3)]},
+      },
+      // Snapshot epoch
+      {
+        title: "slot before epoch 7",
+        slot: computeStartSlotAtEpoch(7) - 1,
+        output: {snapshotSlot: 0, diffSlots: [computeStartSlotAtEpoch(5), computeStartSlotAtEpoch(6)]},
+      },
+      {
+        title: "slot at epoch 7",
+        slot: computeStartSlotAtEpoch(7),
+        output: {snapshotSlot: computeStartSlotAtEpoch(7), diffSlots: []},
+      },
+      {
+        title: "slot after epoch 7",
+        slot: computeStartSlotAtEpoch(7) + 1,
+        output: {snapshotSlot: computeStartSlotAtEpoch(7), diffSlots: []},
+      },
+      // An epoch after first snapshot
+      {
+        title: "slot before epoch 8",
+        slot: computeStartSlotAtEpoch(8) - 1,
+        output: {snapshotSlot: computeStartSlotAtEpoch(7), diffSlots: []},
+      },
+      {
+        title: "slot at epoch 8",
+        slot: computeStartSlotAtEpoch(8),
+        output: {snapshotSlot: computeStartSlotAtEpoch(7), diffSlots: [computeStartSlotAtEpoch(8)]},
+      },
+      {
+        title: "slot after epoch 8",
+        slot: computeStartSlotAtEpoch(8) + 1,
+        output: {snapshotSlot: computeStartSlotAtEpoch(7), diffSlots: [computeStartSlotAtEpoch(8)]},
+      },
+    ];
+
+    it.each(testCases)("$title", ({slot, output}) => {
+      expect(hierarchicalLayers.getArchiveLayers(slot)).toEqual(output);
+    });
+  });
+});