Fixed test setups

consensus/integration/epoch_test.go

@@ -1,6 +1,7 @@
 package integration_test
 
 import (
+	"github.com/onflow/flow-go/state/protocol"
 	"testing"
 	"time"
 
@@ -255,5 +256,26 @@ func withNextEpoch(
 			Map(mapfunc.WithWeight(0))...,
 	).Sort(order.Canonical)
 
+	// update protocol state
+	protocolState := encodableSnapshot.ProtocolState
+	// update protocol state identities since we are in committed phase
+	protocolState.Identities = flow.DynamicIdentityEntryListFromIdentities(encodableSnapshot.Identities)
+	// setup ID has changed, need to update it
+	convertedEpochSetup, _ := protocol.ToEpochSetup(inmem.NewEpoch(*currEpoch))
+	protocolState.CurrentEpochEventIDs.SetupID = convertedEpochSetup.ID()
+	// create next epoch protocol state
+	convertedEpochSetup, _ = protocol.ToEpochSetup(inmem.NewEpoch(*encodableSnapshot.Epochs.Next))
+	convertedEpochCommit, _ := protocol.ToEpochCommit(inmem.NewEpoch(*encodableSnapshot.Epochs.Next))
+	protocolState.NextEpochProtocolState = &flow.ProtocolStateEntry{
+		CurrentEpochEventIDs: flow.EventIDs{
+			SetupID:  convertedEpochSetup.ID(),
+			CommitID: convertedEpochCommit.ID(),
+		},
+		PreviousEpochEventIDs:           protocolState.CurrentEpochEventIDs,
+		Identities:                      flow.DynamicIdentityEntryListFromIdentities(encodableSnapshot.Identities),
+		InvalidStateTransitionAttempted: false,
+		NextEpochProtocolState:          nil,
+	}
+
 	return inmem.SnapshotFromEncodable(encodableSnapshot)
 }

state/protocol/badger/state.go

@@ -195,18 +195,9 @@ func Bootstrap(
 			state.metrics.BlockFinalized(block)
 		}
 
-		rootProtocolState, err := root.ProtocolState()
+		err = state.bootstrapProtocolState(segment, root, protocolState)(tx)
 		if err != nil {
-			return fmt.Errorf("could not get root protocol state: %w", err)
-		}
-		rootProtocolStateEntry := rootProtocolState.Entry()
-		err = protocolState.StoreTx(rootProtocolStateEntry.ID(), rootProtocolStateEntry)(tx)
-		if err != nil {
-			return fmt.Errorf("could not insert root protocol state: %w", err)
-		}
-		err = protocolState.Index(lastFinalized.ID(), rootProtocolStateEntry.ID())(tx)
-		if err != nil {
-			return fmt.Errorf("could not index root protocol state: %w", err)
+			return fmt.Errorf("could not bootstrap protocol state: %w", err)
 		}
 
 		// 7) initialize version beacon
@@ -230,6 +221,29 @@ func Bootstrap(
 	return state, nil
 }
 
+func (state *State) bootstrapProtocolState(segment *flow.SealingSegment, root protocol.Snapshot, protocolState storage.ProtocolState) func(tx *transaction.Tx) error {
+	return func(tx *transaction.Tx) error {
+		rootProtocolState, err := root.ProtocolState()
+		if err != nil {
+			return fmt.Errorf("could not get root protocol state: %w", err)
+		}
+		rootProtocolStateEntry := rootProtocolState.Entry()
+		protocolStateID := rootProtocolStateEntry.ID()
+		err = protocolState.StoreTx(protocolStateID, rootProtocolStateEntry)(tx)
+		if err != nil {
+			return fmt.Errorf("could not insert root protocol state: %w", err)
+		}
+
+		for _, block := range segment.AllBlocks() {
+			err = protocolState.Index(block.ID(), protocolStateID)(tx)
+			if err != nil {
+				return fmt.Errorf("could not index root protocol state: %w", err)
+			}
+		}
+		return nil
+	}
+}
+
 // bootstrapSealingSegment inserts all blocks and associated metadata for the
 // protocol state root snapshot to disk.
 func (state *State) bootstrapSealingSegment(segment *flow.SealingSegment, head *flow.Block, rootSeal *flow.Seal) func(tx *transaction.Tx) error {

state/protocol/inmem/epoch.go

@@ -19,6 +19,10 @@ type Epoch struct {
 
 var _ protocol.Epoch = (*Epoch)(nil)
 
+func NewEpoch(enc EncodableEpoch) Epoch {
+	return Epoch{enc}
+}
+
 func (e Epoch) Encodable() EncodableEpoch {
 	return e.enc
 }