Add tests for ElementAccumulator encoding, StorageProofRoot calculation, and ElementProof update logic

consensus/merkle_test.go

@@ -1,6 +1,11 @@
 package consensus
 
-import "testing"
+import (
+    "bytes"
+    "testing"
+
+    "go.sia.tech/core/types"
+)
 
 func TestElementAccumulatorEncoding(t *testing.T) {
 	for _, test := range []struct {
@@ -37,3 +42,130 @@
 		}
 	}
 }
+
+func TestStorageProofRoot(t *testing.T) {
+    leafHash := types.Hash256{0x01, 0x02, 0x03}
+    validProof := []types.Hash256{
+        {0x0A, 0x0B, 0x0C},
+        {0x0D, 0x0E, 0x0F},
+        {0x10, 0x11, 0x12},
+        {0x13, 0x14, 0x15},
+    }
+    longProof := append(validProof, types.Hash256{0x16, 0x17, 0x18})
+    shortProof := validProof[:2]
+
+    var validWantRoot types.Hash256
+    validWantRoot.UnmarshalText([]byte(`h:a89dbbb545aa4b46696230d104076f06b57a6ab08b2341c3d012bdc13e23eb35`))
+
+    tests := []struct {
+        name      string
+        leafHash  types.Hash256
+        leafIndex uint64
+        filesize  uint64
+        proof     []types.Hash256
+        wantRoot  types.Hash256
+        valid     bool
+    }{
+        {
+            name:      "ValidProof",
+            leafHash:  leafHash,
+            leafIndex: 10,
+            filesize:  829,
+            proof:     validProof,
+            wantRoot:  validWantRoot,
+            valid:     true,
+        },
+        {
+            name:      "TooLongProof",
+            leafHash:  leafHash,
+            leafIndex: 10,
+            filesize:  829,
+            proof:     longProof,
+            wantRoot:  validWantRoot,
+            valid:     false,
+        },
+        {
+            name:      "TooShortProof",
+            leafHash:  leafHash,
+            leafIndex: 10,
+            filesize:  829,
+            proof:     shortProof,
+            wantRoot:  validWantRoot,
+            valid:     false,
+        },
+    }
+
+    for _, tt := range tests {
+        t.Run(tt.name, func(t *testing.T) {
+            gotRoot := storageProofRoot(tt.leafHash, tt.leafIndex, tt.filesize, tt.proof)
+            if tt.valid {
+                if !bytes.Equal(gotRoot[:], tt.wantRoot[:]) {
+                    t.Errorf("%s failed: got %v, want %v", tt.name, gotRoot, tt.wantRoot)
+                }
+            } else {
+                if bytes.Equal(gotRoot[:], tt.wantRoot[:]) {
+                    t.Errorf("%s failed: got a valid root for an invalid proof", tt.name)
+                }
+            }
+        })
+    }
+}
+
+func TestUpdateElementProof(t *testing.T) {
+    tests := []struct {
+        name            string
+        leafIndex       uint64
+        numLeaves       uint64
+        expectPanic     bool
+        expectProofLen  int
+    }{
+        {
+            name:            "EphemeralLeafIndexPanic",
+            leafIndex:       types.EphemeralLeafIndex,
+            numLeaves:       5,
+            expectPanic:     true,
+        },
+        {
+            name:            "LeafIndexOutOfRangePanic",
+            leafIndex:       10,
+            numLeaves:       5,
+            expectPanic:     true,
+        },
+        {
+            name:            "ValidUpdate",
+            leafIndex:       3,
+            numLeaves:       5,
+            expectPanic:     false,
+            expectProofLen:  2,
+        },
+    }
+
+    for _, tt := range tests {
+        t.Run(tt.name, func(t *testing.T) {
+            e := &types.StateElement{
+                LeafIndex:   tt.leafIndex,
+                MerkleProof: make([]types.Hash256, 3),
+            }
+            eru := &elementRevertUpdate{
+                numLeaves: tt.numLeaves,
+                updated:   [64][]elementLeaf{},
+            }
+
+            var didPanic bool
+            func() {
+                defer func() { didPanic = recover() != nil }()
+                eru.updateElementProof(e)
+            }()
+
+            if didPanic != tt.expectPanic {
+                t.Errorf("updateElementProof() didPanic = %v, want %v for %s", didPanic, tt.expectPanic, tt.name)
+            }
+
+            if !tt.expectPanic {
+                if len(e.MerkleProof) != tt.expectProofLen {
+                    t.Errorf("%s: expected MerkleProof length %d, got %d", tt.name, tt.expectProofLen, len(e.MerkleProof))
+                }
+            }
+        })
+    }
+}