fix(dot/digest): verify if next epoch already contains some definition (#2472)

* fix: verify if next epoch already contains some definition

* chore: upgrade mockery version

* chore: remove duplicated code with generics

* chore: improve test comment

Co-authored-by: Quentin McGaw <[email protected]>

* chore: improve test comment

Co-authored-by: Quentin McGaw <[email protected]>

* chore: improve test comment

Co-authored-by: Quentin McGaw <[email protected]>

* chore: improve test comment

Co-authored-by: Quentin McGaw <[email protected]>

* chore: remove the go:generate from the `EpochState`'s comment

Co-authored-by: Quentin McGaw <[email protected]>

* chore: use functions to define mocks in the test cases

* chore: split into two different functions and check for existence before persit

* chore: fix the ci lint warns

* chore: use a more descriptive name

* chore: unexport `Data` field on epoch_test.go

* chore: addressing last nits

* chore: rename to a better name

* chore: fix comments

* chore: addressing comments

* chore: update mockery verison

* chore: update devnet mocks as well

* chore: improve comments

Co-authored-by: Quentin McGaw <[email protected]>

* chore: improve comments

Co-authored-by: Quentin McGaw <[email protected]>

Co-authored-by: Quentin McGaw <[email protected]>

Author: EclesioMeloJunior

devnet/cmd/scale-down-ecs-service/mocks/ecsapi.go

@@ -8,7 +8,6 @@ import (
 	"errors"
 	"fmt"
 
-	"github.com/ChainSafe/gossamer/dot/state"
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/internal/log"
 	"github.com/ChainSafe/gossamer/lib/services"
@@ -17,8 +16,6 @@ import (
 
 var (
 	_ services.Service = &Handler{}
-
-	ErrDefineNextEpoch = errors.New("cannot define next epoch data and config")
 )
 
 // Handler is used to handle consensus messages and relevant authority updates to BABE and GRANDPA
@@ -242,9 +239,14 @@ func (h *Handler) handleBlockFinalisation(ctx context.Context) {
 				continue
 			}
 
-			err := h.persistBABEDigestsForNextEpoch(&info.Header)
+			err := h.epochState.FinalizeBABENextEpochData(&info.Header)
+			if err != nil {
+				h.logger.Errorf("failed to persist babe next epoch data: %s", err)
+			}
+
+			err = h.epochState.FinalizeBABENextConfigData(&info.Header)
 			if err != nil {
-				h.logger.Errorf("failed to store babe next epoch digest: %s", err)
+				h.logger.Errorf("failed to persist babe next epoch config: %s", err)
 			}
 
 			err = h.handleGrandpaChangesOnFinalization(info.Header.Number)
@@ -257,35 +259,6 @@ func (h *Handler) handleBlockFinalisation(ctx context.Context) {
 	}
 }
 
-// persistBABEDigestsForNextEpoch is called only when a block is finalised
-// and defines the correct next epoch data and next config data.
-func (h *Handler) persistBABEDigestsForNextEpoch(finalizedHeader *types.Header) error {
-	currEpoch, err := h.epochState.GetEpochForBlock(finalizedHeader)
-	if err != nil {
-		return fmt.Errorf("cannot get epoch for block %d (%s): %w",
-			finalizedHeader.Number, finalizedHeader.Hash(), err)
-	}
-
-	nextEpoch := currEpoch + 1
-	err = h.epochState.FinalizeBABENextEpochData(nextEpoch)
-	if err != nil && !errors.Is(err, state.ErrEpochNotInMemory) {
-		return fmt.Errorf("cannot finalize babe next epoch data for block number %d (%s): %w",
-			finalizedHeader.Number, finalizedHeader.Hash(), err)
-	}
-
-	err = h.epochState.FinalizeBABENextConfigData(nextEpoch)
-	if err == nil {
-		return nil
-	} else if errors.Is(err, state.ErrEpochNotInMemory) {
-		return fmt.Errorf("%w: %s", ErrDefineNextEpoch, err)
-	}
-
-	// the epoch state does not contains any information about the next epoch
-	return fmt.Errorf("cannot finalize babe next config data for block number %d (%s): %w",
-		finalizedHeader.Number, finalizedHeader.Hash(), err)
-
-}
-
 func (h *Handler) handleGrandpaChangesOnImport(num uint) error {
 	resume := h.grandpaResume
 	if resume != nil && num >= resume.atBlock {

dot/digest/digest.go

@@ -18,6 +18,8 @@ type BlockState interface {
 	FreeFinalisedNotifierChannel(ch chan *types.FinalisationInfo)
 }
 
+//go:generate mockgen -destination=mock_epoch_state_test.go -package $GOPACKAGE . EpochState
+
 // EpochState is the interface for state.EpochState
 type EpochState interface {
 	GetEpochForBlock(header *types.Header) (uint64, error)
@@ -26,8 +28,8 @@ type EpochState interface {
 
 	StoreBABENextEpochData(epoch uint64, hash common.Hash, nextEpochData types.NextEpochData)
 	StoreBABENextConfigData(epoch uint64, hash common.Hash, nextEpochData types.NextConfigData)
-	FinalizeBABENextEpochData(epoch uint64) error
-	FinalizeBABENextConfigData(epoch uint64) error
+	FinalizeBABENextEpochData(finalizedHeader *types.Header) error
+	FinalizeBABENextConfigData(finalizedHeader *types.Header) error
 }
 
 // GrandpaState is the interface for the state.GrandpaState