test(imt.sol): add tests to cover 100% of the lines in the lean imt

aakansha1234 · Mar 22, 2024 · ca07bc4 · ca07bc4
1 parent 70c99d3
commit ca07bc4
Showing 1 changed file with 33 additions and 4 deletions.
diff --git a/packages/imt.sol/test/LeanIMT.ts b/packages/imt.sol/test/LeanIMT.ts
@@ -19,7 +19,7 @@ describe("LeanIMT", () => {
     })
 
     describe("# insert", () => {
-        it("Should not insert a leaf if its value is > SNARK_SCALAR_FIELD", async () => {
+        it("Should not insert a leaf if its value is >= SNARK_SCALAR_FIELD", async () => {
             const transaction = leanIMTTest.insert(SNARK_SCALAR_FIELD)
 
             await expect(transaction).to.be.revertedWithCustomError(leanIMT, "LeafGreaterThanSnarkScalarField")
@@ -65,7 +65,7 @@ describe("LeanIMT", () => {
     })
 
     describe("# insertMany", () => {
-        it("Should not insert a leaf if its value is > SNARK_SCALAR_FIELD", async () => {
+        it("Should not insert a leaf if its value is >= SNARK_SCALAR_FIELD", async () => {
             const transaction = leanIMTTest.insertMany([SNARK_SCALAR_FIELD])
 
             await expect(transaction).to.be.revertedWithCustomError(leanIMT, "LeafGreaterThanSnarkScalarField")
@@ -131,7 +131,7 @@ describe("LeanIMT", () => {
             await expect(transaction).to.be.revertedWithCustomError(leanIMT, "LeafDoesNotExist")
         })
 
-        it("Should not update a leaf if its value is > SNARK_SCALAR_FIELD", async () => {
+        it("Should not update a leaf if its value is >= SNARK_SCALAR_FIELD", async () => {
             const transaction = leanIMTTest.update(2, SNARK_SCALAR_FIELD, [1, 2, 3, 4])
 
             await expect(transaction).to.be.revertedWithCustomError(leanIMT, "LeafGreaterThanSnarkScalarField")
@@ -168,7 +168,7 @@ describe("LeanIMT", () => {
             expect(root).to.equal(jsLeanIMT.root)
         })
 
-        it("Should not update a leaf if the value of at least one sibling node is > SNARK_SCALAR_FIELD", async () => {
+        it("Should not update a leaf if its index is even and the value of at least one sibling node is >= SNARK_SCALAR_FIELD", async () => {
             await leanIMTTest.insert(1)
             await leanIMTTest.insert(2)
 
@@ -184,6 +184,22 @@ describe("LeanIMT", () => {
             await expect(transaction).to.be.revertedWithCustomError(leanIMT, "LeafGreaterThanSnarkScalarField")
         })
 
+        it("Should not update a leaf if its index is odd and the value of at least one sibling node is >= SNARK_SCALAR_FIELD", async () => {
+            await leanIMTTest.insert(1)
+            await leanIMTTest.insert(2)
+
+            jsLeanIMT.insertMany([BigInt(1), BigInt(2)])
+            jsLeanIMT.update(1, BigInt(3))
+
+            const { siblings } = jsLeanIMT.generateProof(1)
+
+            siblings[0] = SNARK_SCALAR_FIELD
+
+            const transaction = leanIMTTest.update(2, 3, siblings)
+
+            await expect(transaction).to.be.revertedWithCustomError(leanIMT, "LeafGreaterThanSnarkScalarField")
+        })
+
         it("Should not update a leaf if the siblings are wrong", async () => {
             await leanIMTTest.insert(1)
             await leanIMTTest.insert(2)
@@ -238,6 +254,19 @@ describe("LeanIMT", () => {
 
             await expect(transaction).to.be.revertedWithCustomError(leanIMT, "LeafCannotBeZero")
         })
+        it("Should not update a leaf if the new value already exists", async () => {
+            await leanIMTTest.insert(1)
+            await leanIMTTest.insert(2)
+
+            jsLeanIMT.insertMany([BigInt(1), BigInt(2)])
+            jsLeanIMT.update(0, BigInt(2))
+
+            const { siblings } = jsLeanIMT.generateProof(0)
+
+            const transaction = leanIMTTest.update(1, 2, siblings)
+
+            await expect(transaction).to.be.revertedWithCustomError(leanIMT, "LeafAlreadyExists")
+        })
     })
 
     describe("# remove", () => {