refactor tests based on PR#2252

NethermindEth · Nov 5, 2024 · d4d7d67 · d4d7d67
1 parent 266604a
commit d4d7d67
Showing 1 changed file with 22 additions and 12 deletions.
diff --git a/rpc/storage_test.go b/rpc/storage_test.go
@@ -159,20 +159,24 @@ func TestStorageProof(t *testing.T) {
 	t.Run("Trie proofs sanity check", func(t *testing.T) {
 		t.Parallel()
 
-		kbs := key.Bytes()
-		kKey := trie.NewKey(251, kbs[:])
-		proof, err := trie.GetProof(&kKey, tempTrie)
+		kKey := tempTrie.FeltToKey(key)
+		proof := trie.NewProofSet()
+		err := tempTrie.Prove(key, proof)
 		require.NoError(t, err)
 		root, err := tempTrie.Root()
 		require.NoError(t, err)
-		require.True(t, trie.VerifyProof(root, &kKey, value, proof, tempTrie.HashFunc()))
+		leaf, err := trie.VerifyProof(root, &kKey, proof, tempTrie.HashFunc())
+		require.NoError(t, err)
+		require.Equal(t, leaf, value)
 
 		// non-membership test
-		kbs = noSuchKey.Bytes()
-		kKey = trie.NewKey(251, kbs[:])
-		proof, err = trie.GetProof(&kKey, tempTrie)
+		kKey = tempTrie.FeltToKey(noSuchKey)
+		proof = trie.NewProofSet()
+		err = tempTrie.Prove(key, proof)
+		require.NoError(t, err)
+		leaf, err = trie.VerifyProof(root, &kKey, proof, tempTrie.HashFunc())
 		require.NoError(t, err)
-		require.True(t, trie.VerifyProof(root, &kKey, nil, proof, tempTrie.HashFunc()))
+		require.Equal(t, felt.Zero, *leaf)
 	})
 	t.Run("global roots are filled", func(t *testing.T) {
 		t.Parallel()
@@ -798,7 +802,6 @@ func TestVerifyPathfinderResponse(t *testing.T) {
 	t.Run("second contract proof verification", func(t *testing.T) {
 		t.Parallel()
 
-		t.Skip("verification issue on the length=1 edge node")
 		secondContractAddr := utils.HexToFelt(t, "0x5fbaa249500be29fee38fdd90a7a2651a8d3935c14167570f6863f563d838f0")
 		secondContractLeaf := utils.HexToFelt(t, "0x25790175fe1fbeed47cbf510a41fba8676bea20a0c8888d4b9090b8f5cf19b8")
 		verifyIf(t,
@@ -840,14 +843,21 @@ func verifyIf(
 ) {
 	t.Helper()
 
-	pnodes := []trie.ProofNode{}
+	proofSet := trie.NewProofSet()
 	for _, hn := range proof {
-		pnodes = append(pnodes, hn.Node.AsProofNode())
+		proofSet.Put(*hn.Hash, hn.Node.AsProofNode())
 	}
 
 	kbs := key.Bytes()
 	kkey := trie.NewKey(251, kbs[:])
-	require.True(t, trie.VerifyProof(root, &kkey, value, pnodes, hashF))
+	leaf, err := trie.VerifyProof(root, &kkey, proofSet, hashF)
+	require.NoError(t, err)
+
+	// non-membership test
+	if value == nil {
+		value = felt.Zero.Clone()
+	}
+	require.Equal(t, leaf, value)
 }
 
 func verifyGlobalStateRoot(t *testing.T, globalStateRoot, classRoot, storageRoot *felt.Felt) {