From 8757ad44995f72bdcc1b872948e5443af1659dfe Mon Sep 17 00:00:00 2001
From: Zsolt Felfoldi <zsfelfoldi@gmail.com>
Date: Fri, 17 Jan 2025 01:20:38 +0100
Subject: [PATCH] triedb/pathdb: add rawStorageKey check to encode/decode unit
 tests

---
 triedb/pathdb/history_test.go | 11 ++++++++---
 triedb/pathdb/states_test.go  | 20 ++++++++++++++++++--
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/triedb/pathdb/history_test.go b/triedb/pathdb/history_test.go
index 2e1aabbe4e..953f023530 100644
--- a/triedb/pathdb/history_test.go
+++ b/triedb/pathdb/history_test.go
@@ -49,9 +49,9 @@ func randomStateSet(n int) (map[common.Address][]byte, map[common.Address]map[co
 	return accounts, storages
 }
 
-func makeHistory() *history {
+func makeHistory(rawStorageKey bool) *history {
 	accounts, storages := randomStateSet(3)
-	return newHistory(testrand.Hash(), types.EmptyRootHash, 0, accounts, storages, false)
+	return newHistory(testrand.Hash(), types.EmptyRootHash, 0, accounts, storages, rawStorageKey)
 }
 
 func makeHistories(n int) []*history {
@@ -70,10 +70,15 @@ func makeHistories(n int) []*history {
 }
 
 func TestEncodeDecodeHistory(t *testing.T) {
+	testEncodeDecodeHistory(t, false)
+	testEncodeDecodeHistory(t, true)
+}
+
+func testEncodeDecodeHistory(t *testing.T, rawStorageKey bool) {
 	var (
 		m   meta
 		dec history
-		obj = makeHistory()
+		obj = makeHistory(rawStorageKey)
 	)
 	// check if meta data can be correctly encode/decode
 	blob := obj.meta.encode()
diff --git a/triedb/pathdb/states_test.go b/triedb/pathdb/states_test.go
index 1860a342b3..30eb6ad6c8 100644
--- a/triedb/pathdb/states_test.go
+++ b/triedb/pathdb/states_test.go
@@ -290,6 +290,11 @@ func TestStateRevertStorageNullMarker(t *testing.T) {
 }
 
 func TestStatesEncode(t *testing.T) {
+	testStatesEncode(t, false)
+	testStatesEncode(t, true)
+}
+
+func testStatesEncode(t *testing.T, rawStorageKey bool) {
 	s := newStates(
 		map[common.Hash][]byte{
 			{0x1}: {0x1},
@@ -299,7 +304,7 @@ func TestStatesEncode(t *testing.T) {
 				common.Hash{0x1}: {0x1},
 			},
 		},
-		false,
+		rawStorageKey,
 	)
 	buf := bytes.NewBuffer(nil)
 	if err := s.encode(buf); err != nil {
@@ -315,9 +320,17 @@ func TestStatesEncode(t *testing.T) {
 	if !reflect.DeepEqual(s.storageData, dec.storageData) {
 		t.Fatal("Unexpected storage data")
 	}
+	if s.rawStorageKey != dec.rawStorageKey {
+		t.Fatal("Unexpected rawStorageKey flag")
+	}
 }
 
 func TestStateWithOriginEncode(t *testing.T) {
+	testStateWithOriginEncode(t, false)
+	testStateWithOriginEncode(t, true)
+}
+
+func testStateWithOriginEncode(t *testing.T, rawStorageKey bool) {
 	s := NewStateSetWithOrigin(
 		map[common.Hash][]byte{
 			{0x1}: {0x1},
@@ -335,7 +348,7 @@ func TestStateWithOriginEncode(t *testing.T) {
 				common.Hash{0x1}: {0x1},
 			},
 		},
-		false,
+		rawStorageKey,
 	)
 	buf := bytes.NewBuffer(nil)
 	if err := s.encode(buf); err != nil {
@@ -357,6 +370,9 @@ func TestStateWithOriginEncode(t *testing.T) {
 	if !reflect.DeepEqual(s.storageOrigin, dec.storageOrigin) {
 		t.Fatal("Unexpected storage origin data")
 	}
+	if s.rawStorageKey != dec.rawStorageKey {
+		t.Fatal("Unexpected rawStorageKey flag")
+	}
 }
 
 func TestStateSizeTracking(t *testing.T) {