Dependency inject to NewBlockStateFromGenesis

ChainSafe · Dec 6, 2022 · 8f14659 · 8f14659
1 parent 92aed5e
commit 8f14659
Show file tree

Hide file tree

Showing 7 changed files with 34 additions and 12 deletions.
diff --git a/dot/rpc/modules/dev_integration_test.go b/dot/rpc/modules/dev_integration_test.go
@@ -9,6 +9,7 @@ import (
 	"encoding/binary"
 	"testing"
 
+	"github.com/ChainSafe/chaindb"
 	"github.com/ChainSafe/gossamer/dot/state"
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/lib/babe"
@@ -43,7 +44,12 @@ func newState(t *testing.T) (*state.BlockState, *state.EpochState) {
 	_, genesisTrie, genesisHeader := newTestGenesisWithTrieAndHeader(t)
 	tries := state.NewTries()
 	tries.SetTrie(&genesisTrie)
-	bs, err := state.NewBlockStateFromGenesis(db, tries, &genesisHeader, telemetryMock)
+
+	baseState := state.NewBaseState(db)
+	const blockPrefix = "block"
+	blockStateDatabase := chaindb.NewTable(db, blockPrefix)
+	bs, err := state.NewBlockStateFromGenesis(blockStateDatabase,
+		baseState, tries, &genesisHeader, telemetryMock)
 	require.NoError(t, err)
 	es, err := state.NewEpochStateFromGenesis(db, bs, genesisBABEConfig)
 	require.NoError(t, err)

diff --git a/dot/state/block.go b/dot/state/block.go
@@ -25,6 +25,8 @@ import (
 )
 
 const (
+	// blockPrefix is the prefix to use for the block state
+	// key value database.
 	blockPrefix = "block"
 )
 
@@ -101,13 +103,13 @@ func NewBlockState(db BlockStateDatabase, baseState *BaseState,
 }
 
 // NewBlockStateFromGenesis initialises a BlockState from a genesis header,
-// saving it to the database located at basePath
-func NewBlockStateFromGenesis(db *chaindb.BadgerDB, trs *Tries, header *types.Header,
-	telemetryMailer Telemetry) (*BlockState, error) {
+// saving it to the database.
+func NewBlockStateFromGenesis(db BlockStateDatabase, baseState *BaseState,
+	trs *Tries, header *types.Header, telemetryMailer Telemetry) (*BlockState, error) {
 	bs := &BlockState{
 		bt:                         blocktree.NewBlockTreeFromRoot(header),
-		baseState:                  NewBaseState(db),
-		db:                         chaindb.NewTable(db, blockPrefix),
+		baseState:                  baseState,
+		db:                         db,
 		unfinalisedBlocks:          newHashToBlockMap(),
 		tries:                      trs,
 		imported:                   make(map[chan *types.Block]struct{}),

diff --git a/dot/state/block_race_test.go b/dot/state/block_race_test.go
@@ -35,7 +35,11 @@ func TestConcurrencySetHeader(t *testing.T) {
 		go func(index int) {
 			defer pend.Done()
 
-			bs, err := NewBlockStateFromGenesis(dbs[index], tries, testGenesisHeader, telemetryMock)
+			db := dbs[index]
+			baseState := NewBaseState(db)
+			blockStateDatabase := chaindb.NewTable(db, blockPrefix)
+			bs, err := NewBlockStateFromGenesis(blockStateDatabase,
+				baseState, tries, testGenesisHeader, telemetryMock)
 			require.NoError(t, err)
 
 			header := &types.Header{

diff --git a/dot/state/block_test.go b/dot/state/block_test.go
@@ -7,6 +7,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/ChainSafe/chaindb"
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/ChainSafe/gossamer/lib/trie"
@@ -32,7 +33,10 @@ func newTestBlockState(t *testing.T, tries *Tries) *BlockState {
 	db := NewInMemoryDB(t)
 	header := testGenesisHeader
 
-	bs, err := NewBlockStateFromGenesis(db, tries, header, telemetryMock)
+	baseState := NewBaseState(db)
+	blockStateDatabase := chaindb.NewTable(db, blockPrefix)
+	bs, err := NewBlockStateFromGenesis(blockStateDatabase,
+		baseState, tries, header, telemetryMock)
 	require.NoError(t, err)
 
 	// loads in-memory tries with genesis state root, should be deleted

diff --git a/dot/state/grandpa_test.go b/dot/state/grandpa_test.go
@@ -129,7 +129,9 @@ func testBlockState(t *testing.T, db *chaindb.BadgerDB) *BlockState {
 	telemetryMock.EXPECT().SendMessage(gomock.Any()).AnyTimes()
 	header := testGenesisHeader
 
-	bs, err := NewBlockStateFromGenesis(db, newTriesEmpty(), header, telemetryMock)
+	baseState := NewBaseState(db)
+	blockStateDatabase := chaindb.NewTable(db, blockPrefix)
+	bs, err := NewBlockStateFromGenesis(blockStateDatabase, baseState, newTriesEmpty(), header, telemetryMock)
 	require.NoError(t, err)
 
 	// loads in-memory tries with genesis state root, should be deleted

diff --git a/dot/state/initialize.go b/dot/state/initialize.go
@@ -66,8 +66,9 @@ func (s *Service) Initialise(gen *genesis.Genesis, header *types.Header, t *trie
 	tries := NewTries()
 	tries.SetTrie(t)
 
-	// create block state from genesis block
-	blockState, err := NewBlockStateFromGenesis(db, tries, header, s.Telemetry)
+	baseState := NewBaseState(db)
+	blockStateDatabase := chaindb.NewTable(db, blockPrefix)
+	blockState, err := NewBlockStateFromGenesis(blockStateDatabase, baseState, tries, header, s.Telemetry)
 	if err != nil {
 		return fmt.Errorf("failed to create block state from genesis: %s", err)
 	}

diff --git a/dot/state/storage_test.go b/dot/state/storage_test.go
@@ -7,6 +7,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/ChainSafe/chaindb"
 	"github.com/ChainSafe/gossamer/dot/state/pruner"
 	"github.com/ChainSafe/gossamer/dot/telemetry"
 	"github.com/ChainSafe/gossamer/dot/types"
@@ -194,7 +195,9 @@ func TestGetStorageChildAndGetStorageFromChild(t *testing.T) {
 
 	tries := newTriesEmpty()
 
-	blockState, err := NewBlockStateFromGenesis(db, tries, &genHeader, telemetryMock)
+	baseState := NewBaseState(db)
+	blockStateDatabase := chaindb.NewTable(db, blockPrefix)
+	blockState, err := NewBlockStateFromGenesis(blockStateDatabase, baseState, tries, &genHeader, telemetryMock)
 	require.NoError(t, err)
 
 	storage, err := NewStorageState(db, blockState, tries, pruner.Config{})