From 958f5343613d2a2853d499c2783715133cc66dcb Mon Sep 17 00:00:00 2001
From: Bryan White <bryanchriswhite@gmail.com>
Date: Tue, 16 Jul 2024 09:25:05 +0200
Subject: [PATCH] test: rewrite MerkleSumRoot tests

---
 root_test.go | 100 ++++++++++++++++++++++++++++++---------------------
 1 file changed, 59 insertions(+), 41 deletions(-)

diff --git a/root_test.go b/root_test.go
index b417ff5..ea37fdf 100644
--- a/root_test.go
+++ b/root_test.go
@@ -13,64 +13,82 @@ import (
 	"github.com/pokt-network/smt/kvstore/simplemap"
 )
 
-func TestMerkleRoot_TrieTypes(t *testing.T) {
+func TestMerkleSumRoot_SumAndCountSuccess(t *testing.T) {
 	tests := []struct {
-		desc          string
-		sumTree       bool
-		hasher        hash.Hash
-		expectedPanic string
+		desc   string
+		hasher hash.Hash
 	}{
 		{
-			desc:          "successfully: gets sum of sha256 hasher SMST",
-			sumTree:       true,
-			hasher:        sha256.New(),
-			expectedPanic: "",
+			desc:   "sha256 hasher",
+			hasher: sha256.New(),
 		},
 		{
-			desc:          "successfully: gets sum of sha512 hasher SMST",
-			sumTree:       true,
-			hasher:        sha512.New(),
-			expectedPanic: "",
+			desc:   "sha512 hasher",
+			hasher: sha512.New(),
 		},
+	}
+
+	nodeStore := simplemap.NewSimpleMap()
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			t.Cleanup(func() {
+				require.NoError(t, nodeStore.ClearAll())
+			})
+			trie := smt.NewSparseMerkleSumTrie(nodeStore, test.hasher)
+			for i := uint64(0); i < 10; i++ {
+				require.NoError(t, trie.Update([]byte(fmt.Sprintf("key%d", i)), []byte(fmt.Sprintf("value%d", i)), i))
+			}
+
+			sum, sumErr := trie.Sum()
+			require.NoError(t, sumErr)
+
+			count, countErr := trie.Count()
+			require.NoError(t, countErr)
+
+			require.EqualValues(t, uint64(45), sum)
+			require.EqualValues(t, uint64(10), count)
+		})
+	}
+}
+
+func TestMekleRoot_SumAndCountError(t *testing.T) {
+	tests := []struct {
+		desc   string
+		hasher hash.Hash
+	}{
 		{
-			desc:          "failure: panics for sha256 hasher SMT",
-			sumTree:       false,
-			hasher:        sha256.New(),
-			expectedPanic: "roo#sum: not a merkle sum trie",
+			desc:   "sha256 hasher",
+			hasher: sha256.New(),
 		},
 		{
-			desc:          "failure: panics for sha512 hasher SMT",
-			sumTree:       false,
-			hasher:        sha512.New(),
-			expectedPanic: "roo#sum: not a merkle sum trie",
+			desc:   "sha512 hasher",
+			hasher: sha512.New(),
 		},
 	}
 
 	nodeStore := simplemap.NewSimpleMap()
-	for _, tt := range tests {
-		tt := tt
-		t.Run(tt.desc, func(t *testing.T) {
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
 			t.Cleanup(func() {
 				require.NoError(t, nodeStore.ClearAll())
 			})
-			if tt.sumTree {
-				trie := smt.NewSparseMerkleSumTrie(nodeStore, tt.hasher)
-				for i := uint64(0); i < 10; i++ {
-					require.NoError(t, trie.Update([]byte(fmt.Sprintf("key%d", i)), []byte(fmt.Sprintf("value%d", i)), i))
-				}
-				require.NotNil(t, trie.MustSum())
-				require.EqualValues(t, 45, trie.MustSum())
-				require.EqualValues(t, 10, trie.MustCount())
-
-				return
-			}
-			trie := smt.NewSparseMerkleTrie(nodeStore, tt.hasher)
-			for i := 0; i < 10; i++ {
-				require.NoError(t, trie.Update([]byte(fmt.Sprintf("key%d", i)), []byte(fmt.Sprintf("value%d", i))))
-			}
-			if panicStr := recover(); panicStr != nil {
-				require.Equal(t, tt.expectedPanic, panicStr)
+			trie := smt.NewSparseMerkleSumTrie(nodeStore, test.hasher)
+			for i := uint64(0); i < 10; i++ {
+				require.NoError(t, trie.Update([]byte(fmt.Sprintf("key%d", i)), []byte(fmt.Sprintf("value%d", i)), i))
 			}
+
+			root := trie.Root()
+
+			// Mangle the root bytes.
+			root = root[:len(root)-1]
+
+			sum, sumErr := root.Sum()
+			require.Error(t, sumErr)
+			require.Equal(t, uint64(0), sum)
+
+			count, countErr := root.Count()
+			require.Error(t, countErr)
+			require.Equal(t, uint64(0), count)
 		})
 	}
 }