From 8f33fc6f7f1b59e49ffc9e7eeda7621272a5c37d Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Thu, 14 Apr 2022 10:41:11 +0000
Subject: [PATCH] chore(lib/babe): epoch error wrapping fixed up

---
 dot/state/epoch.go | 4 +++-
 lib/babe/babe.go   | 6 +++---
 lib/babe/epoch.go  | 2 +-
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/dot/state/epoch.go b/dot/state/epoch.go
index cdc3fba584..969c6360e3 100644
--- a/dot/state/epoch.go
+++ b/dot/state/epoch.go
@@ -323,7 +323,9 @@ func (s *EpochState) HasEpochData(epoch uint64) (bool, error) {
 		return has, nil
 	}
 
-	if !errors.Is(chaindb.ErrKeyNotFound, err) {
+	// we can have `has == false` and `err == nil`
+	// so ensure the error is not nil in the condition below.
+	if err != nil && !errors.Is(chaindb.ErrKeyNotFound, err) {
 		return false, fmt.Errorf("cannot check database for epoch key %d: %w", epoch, err)
 	}
 
diff --git a/lib/babe/babe.go b/lib/babe/babe.go
index 9a3807c64c..d0dd5fa4e7 100644
--- a/lib/babe/babe.go
+++ b/lib/babe/babe.go
@@ -399,7 +399,7 @@ func (b *Service) initiate() {
 func (b *Service) initiateAndGetEpochHandler(epoch uint64) (*epochHandler, error) {
 	epochData, err := b.initiateEpoch(epoch)
 	if err != nil {
-		return nil, fmt.Errorf("failed to initiate epoch %d: %w", epoch, err)
+		return nil, fmt.Errorf("failed to initiate epoch: %w", err)
 	}
 
 	logger.Debugf("initiated epoch with threshold %s, randomness 0x%x and authorities %v",
@@ -430,7 +430,7 @@ func (b *Service) runEngine() error {
 		if errors.Is(err, errServicePaused) || errors.Is(err, context.Canceled) {
 			return nil
 		} else if err != nil {
-			return err
+			return fmt.Errorf("cannot handle epoch: %w", err)
 		}
 
 		epoch = next
@@ -442,7 +442,7 @@ func (b *Service) handleEpoch(epoch uint64) (next uint64, err error) {
 	defer cancel()
 	b.epochHandler, err = b.initiateAndGetEpochHandler(epoch)
 	if err != nil {
-		return 0, fmt.Errorf("cannot initiate and get epoch handler for epoch %d: %w", epoch, err)
+		return 0, fmt.Errorf("cannot initiate and get epoch handler: %w", err)
 	}
 
 	// get start slot for current epoch
diff --git a/lib/babe/epoch.go b/lib/babe/epoch.go
index 379e53be1c..1e524094fb 100644
--- a/lib/babe/epoch.go
+++ b/lib/babe/epoch.go
@@ -92,7 +92,7 @@ func (b *Service) getEpochDataAndStartSlot(epoch uint64) (*epochData, uint64, er
 
 	has, err := b.epochState.HasEpochData(epoch)
 	if err != nil {
-		return nil, 0, fmt.Errorf("cannot check for epoch data for epoch %d: %w", epoch, err)
+		return nil, 0, fmt.Errorf("cannot check epoch state: %w", err)
 	}
 
 	if !has {