From ae4e67de19d5ffabf1e65e7b9c5b077b3e4b4bbf Mon Sep 17 00:00:00 2001
From: Albert <albert@rss3.io>
Date: Tue, 16 Jan 2024 17:25:46 +0800
Subject: [PATCH] feat: add test for statedb AddBalance/SubBalance

---
 core/state/statedb_test.go | 72 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 72 insertions(+)

diff --git a/core/state/statedb_test.go b/core/state/statedb_test.go
index ad829a0c8..83ba1f92e 100644
--- a/core/state/statedb_test.go
+++ b/core/state/statedb_test.go
@@ -1191,3 +1191,75 @@ func TestDeleteStorage(t *testing.T) {
 		t.Fatalf("difference found:\nfast: %v\nslow: %v\n", fastRes, slowRes)
 	}
 }
+
+func TestAddBalance(t *testing.T) {
+	// Create an empty state database
+	var (
+		db  = rawdb.NewMemoryDatabase()
+		tdb = trie.NewDatabase(db, nil)
+	)
+	state, _ := New(types.EmptyRootHash, NewDatabaseWithNodeDB(db, tdb), nil)
+
+	// add balance
+	for i := byte(0); i < 255; i++ {
+		addr := common.BytesToAddress([]byte{i})
+		state.AddBalance(addr, big.NewInt(int64(11*i)))
+	}
+	// Write modifications to trie.
+	root := state.IntermediateRoot(false)
+	if err := tdb.Commit(root, false); err != nil {
+		t.Errorf("can not commit trie %v to persistent database", root.Hex())
+	}
+
+	// check balance
+	for i := byte(0); i < 255; i++ {
+		addr := common.BytesToAddress([]byte{i})
+		balance := state.GetBalance(addr)
+		expectedBalance := big.NewInt(int64(11 * i))
+		if balance.Cmp(expectedBalance) != 0 {
+			t.Errorf("account %d: balance mismatch: have %v, want %v", i, balance, expectedBalance)
+		}
+	}
+}
+
+func TestSubBalance(t *testing.T) {
+	// Create an empty state database
+	var (
+		db  = rawdb.NewMemoryDatabase()
+		tdb = trie.NewDatabase(db, nil)
+	)
+	state, _ := New(types.EmptyRootHash, NewDatabaseWithNodeDB(db, tdb), nil)
+
+	// set balance
+	for i := byte(0); i < 5; i++ {
+		addr := common.BytesToAddress([]byte{i})
+		balance := big.NewInt(int64(11 * i))
+		state.SetBalance(addr, balance)
+	}
+
+	// sub balance
+	for i := byte(0); i < 5; i++ {
+		addr := common.BytesToAddress([]byte{i})
+		state.SubBalance(addr, big.NewInt(int64(1)))
+	}
+
+	// Write modifications to trie.
+	root := state.IntermediateRoot(false)
+	if err := tdb.Commit(root, false); err != nil {
+		t.Errorf("can not commit trie %v to persistent database", root.Hex())
+	}
+
+	// check balance
+	for i := byte(0); i < 5; i++ {
+		addr := common.BytesToAddress([]byte{i})
+		balance := state.GetBalance(addr)
+		expectedBalance := big.NewInt(int64(11*i) - 1)
+		if expectedBalance.Cmp(big.NewInt(int64(-1))) == 0 {
+			// for account 0
+			expectedBalance = big.NewInt(int64(1))
+		}
+		if balance.Cmp(expectedBalance) != 0 {
+			t.Errorf("account %d: balance mismatch: have %v, want %v", i, balance, expectedBalance)
+		}
+	}
+}