exp/lighthorizon: Improve trie tests to avoid raw comparisons/outputs. (

#4373)
stellar · May 12, 2022 · e3397e5 · e3397e5
1 parent 48e51e9
commit e3397e5
Showing 1 changed file with 35 additions and 58 deletions.
diff --git a/exp/lighthorizon/index/trie_test.go b/exp/lighthorizon/index/trie_test.go
@@ -6,8 +6,10 @@ import (
 	"encoding/hex"
 	"encoding/json"
 	"math/rand"
+	"strconv"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -21,7 +23,7 @@ func randomTrie(t *testing.T, index *TrieIndex) (*TrieIndex, map[string]uint32)
 		ledger := uint32(rand.Int63())
 		hashBytes := make([]byte, 32)
 		if _, err := rand.Read(hashBytes); err != nil {
-			t.Error(err.Error())
+			assert.NoError(t, err)
 		}
 		hash := hex.EncodeToString(hashBytes)
 
@@ -39,14 +41,10 @@ func TestTrieIndex(t *testing.T) {
 
 		for key, expected := range inserts {
 			value, ok := index.Get([]byte(key))
-			if !ok {
-				t.Errorf("Key not found: %s", key)
-			} else {
-				ledger := binary.BigEndian.Uint32(value)
-				if ledger != expected {
-					t.Errorf("Key %s found: %v, expected: %v", key, ledger, expected)
-				}
-			}
+			require.Truef(t, ok, "Key not found: %s", key)
+			ledger := binary.BigEndian.Uint32(value)
+			assert.Equalf(t, expected, ledger,
+				"Key %s found: %v, expected: %v", key, ledger, expected)
 		}
 	}
 }
@@ -56,19 +54,16 @@ func TestTrieIndexUpsertBasic(t *testing.T) {
 
 	key := "key"
 	prev, ok := index.Upsert([]byte(key), []byte("a"))
-	if ok || prev != nil {
-		t.Errorf("Unexpected previous value: %q, expected: nil", string(prev))
-	}
+	assert.Nil(t, prev)
+	assert.Falsef(t, ok, "expected nil, got prev: %q", string(prev))
 
 	prev, ok = index.Upsert([]byte(key), []byte("b"))
-	if !ok || string(prev) != "a" {
-		t.Errorf("Unexpected previous value: %q, expected: a", string(prev))
-	}
+	assert.Equal(t, "a", string(prev))
+	assert.Truef(t, ok, "expected 'a', got prev: %q", string(prev))
 
 	prev, ok = index.Upsert([]byte(key), []byte("c"))
-	if !ok || string(prev) != "b" {
-		t.Errorf("Unexpected previous value: %q, expected: b", string(prev))
-	}
+	assert.Equal(t, "b", string(prev))
+	assert.Truef(t, ok, "expected 'b', got prev: %q", string(prev))
 }
 
 func TestTrieIndexSuffixes(t *testing.T) {
@@ -101,36 +96,28 @@ func TestTrieIndexSuffixes(t *testing.T) {
 
 func TestTrieIndexSerialization(t *testing.T) {
 	for i := 0; i < 10_000; i++ {
-		index, inserts := randomTrie(t, nil)
+		t.Run(strconv.FormatInt(int64(i), 10), func(t *testing.T) {
+			index, inserts := randomTrie(t, nil)
 
-		// Round-trip it to serialization and back
-		buf := &bytes.Buffer{}
-		nWritten, err := index.WriteTo(buf)
-		if err != nil {
-			t.Error(err.Error())
-		}
+			// Round-trip it to serialization and back
+			buf := &bytes.Buffer{}
+			nWritten, err := index.WriteTo(buf)
+			assert.NoError(t, err)
 
-		read := &TrieIndex{}
-		nRead, err := read.ReadFrom(buf)
-		if err != nil {
-			t.Error(err.Error())
-		}
+			read := &TrieIndex{}
+			nRead, err := read.ReadFrom(buf)
+			assert.NoError(t, err)
 
-		if nWritten != nRead {
-			t.Errorf("Wrote %d bytes, but read %d bytes", nWritten, nRead)
-		}
+			assert.Equal(t, nWritten, nRead, "read more or less than we wrote")
+
+			for key, expected := range inserts {
+				value, ok := read.Get([]byte(key))
+				require.Truef(t, ok, "Key not found: %s", key)
 
-		for key, expected := range inserts {
-			value, ok := read.Get([]byte(key))
-			if !ok {
-				t.Errorf("Key not found: %s", key)
-			} else {
 				ledger := binary.BigEndian.Uint32(value)
-				if ledger != expected {
-					t.Errorf("Key %s found: %v, expected: %v", key, ledger, expected)
-				}
+				assert.Equal(t, expected, ledger, "for key %s", key)
 			}
-		}
+		})
 	}
 }
 
@@ -294,27 +281,17 @@ func TestTrieIndexMerge(t *testing.T) {
 		// Should still have all the A keys
 		for key, expected := range aInserts {
 			value, ok := a.Get([]byte(key))
-			if !ok {
-				t.Errorf("Key not found: %s", key)
-			} else {
-				ledger := binary.BigEndian.Uint32(value)
-				if ledger != expected {
-					t.Errorf("Key %s found: %v, expected: %v", key, ledger, expected)
-				}
-			}
+			require.Truef(t, ok, "Key not found: %s", key)
+			ledger := binary.BigEndian.Uint32(value)
+			assert.Equalf(t, expected, ledger, "Key %s found", key)
 		}
 
 		// Should now also have all the B keys
 		for key, expected := range bInserts {
 			value, ok := a.Get([]byte(key))
-			if !ok {
-				t.Errorf("Key not found: %s", key)
-			} else {
-				ledger := binary.BigEndian.Uint32(value)
-				if ledger != expected {
-					t.Errorf("Key %s found: %v, expected: %v", key, ledger, expected)
-				}
-			}
+			require.Truef(t, ok, "Key not found: %s", key)
+			ledger := binary.BigEndian.Uint32(value)
+			assert.Equalf(t, expected, ledger, "Key %s found", key)
 		}
 	}
 }