From 4d22f8dc0dab634ad30163513258f731716bd1f5 Mon Sep 17 00:00:00 2001
From: Dang Nhat Trinh <trinhdn@coin98.finance>
Date: Tue, 18 Jul 2023 16:04:29 +0700
Subject: [PATCH] Fix unit tests

---
 core/bench_test.go | 12 ++++++++++--
 core/blockchain.go |  8 ++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/core/bench_test.go b/core/bench_test.go
index 137b57f031..cef95625c6 100644
--- a/core/bench_test.go
+++ b/core/bench_test.go
@@ -18,7 +18,6 @@ package core
 
 import (
 	"crypto/ecdsa"
-	"github.com/tomochain/tomochain/core/rawdb"
 	"io/ioutil"
 	"math/big"
 	"os"
@@ -27,6 +26,7 @@ import (
 	"github.com/tomochain/tomochain/common"
 	"github.com/tomochain/tomochain/common/math"
 	"github.com/tomochain/tomochain/consensus/ethash"
+	"github.com/tomochain/tomochain/core/rawdb"
 	"github.com/tomochain/tomochain/core/types"
 	"github.com/tomochain/tomochain/core/vm"
 	"github.com/tomochain/tomochain/crypto"
@@ -238,10 +238,16 @@ func makeChainForBench(db ethdb.Database, full bool, count uint64) {
 		WriteHeader(db, header)
 		WriteCanonicalHash(db, hash, n)
 		WriteTd(db, hash, n, big.NewInt(int64(n+1)))
+		if n == 0 {
+			WriteChainConfig(db, hash, params.AllEthashProtocolChanges)
+		}
+		WriteHeadHeaderHash(db, hash)
+
 		if full || n == 0 {
 			block := types.NewBlockWithHeader(header)
 			WriteBody(db, hash, n, block.Body())
 			WriteBlockReceipts(db, hash, n, nil)
+			WriteHeadBlockHash(db, hash)
 		}
 	}
 }
@@ -275,6 +281,8 @@ func benchReadChain(b *testing.B, full bool, count uint64) {
 	}
 	makeChainForBench(db, full, count)
 	db.Close()
+	cacheConfig := defaultCacheConfig
+	cacheConfig.Disabled = true
 
 	b.ReportAllocs()
 	b.ResetTimer()
@@ -284,7 +292,7 @@ func benchReadChain(b *testing.B, full bool, count uint64) {
 		if err != nil {
 			b.Fatalf("error opening database at %v: %v", dir, err)
 		}
-		chain, err := NewBlockChain(db, nil, params.TestChainConfig, ethash.NewFaker(), vm.Config{})
+		chain, err := NewBlockChain(db, cacheConfig, params.TestChainConfig, ethash.NewFaker(), vm.Config{})
 		if err != nil {
 			b.Fatalf("error creating chain: %v", err)
 		}
diff --git a/core/blockchain.go b/core/blockchain.go
index f763189be7..18b1521ff5 100644
--- a/core/blockchain.go
+++ b/core/blockchain.go
@@ -82,6 +82,14 @@ type CacheConfig struct {
 	TrieNodeLimit int           // Memory limit (MB) at which to flush the current in-memory trie to disk
 	TrieTimeLimit time.Duration // Time limit after which to flush the current in-memory trie to disk
 }
+
+// defaultCacheConfig are the default caching values if none are specified by the
+// user (also used during testing).
+var defaultCacheConfig = &CacheConfig{
+	TrieNodeLimit: 256,
+	TrieTimeLimit: 5 * time.Minute,
+}
+
 type ResultProcessBlock struct {
 	logs         []*types.Log
 	receipts     []*types.Receipt