From bae9c4400b6756d3686772b7c0b2ef80f91530c8 Mon Sep 17 00:00:00 2001
From: Manav Darji <manavdarji.india@gmail.com>
Date: Thu, 12 May 2022 14:49:48 +0530
Subject: [PATCH] add unit tests for IsValidChain function

---
 eth/downloader/whitelist/service.go      |  2 +-
 eth/downloader/whitelist/service_test.go | 64 ++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/eth/downloader/whitelist/service.go b/eth/downloader/whitelist/service.go
index d79284bfad..7889429c87 100644
--- a/eth/downloader/whitelist/service.go
+++ b/eth/downloader/whitelist/service.go
@@ -38,7 +38,7 @@ func (w *Service) IsValidChain(remoteHeader *types.Header, fetchHeadersByNumber
 	// Check for availaibility of the last checkpointed block.
 	// This can be also be empty if our heimdall is not responsing
 	// or we're running without it.
-	if len(w.checkpointWhitelist) <= 0 {
+	if len(w.checkpointWhitelist) == 0 {
 		// worst case, we don't have the checkpoints in memory
 		return true, nil
 	}
diff --git a/eth/downloader/whitelist/service_test.go b/eth/downloader/whitelist/service_test.go
index dbd3f2f92e..db17af8ba8 100644
--- a/eth/downloader/whitelist/service_test.go
+++ b/eth/downloader/whitelist/service_test.go
@@ -1,9 +1,12 @@
 package whitelist
 
 import (
+	"errors"
+	"math/big"
 	"testing"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/types"
 	"gotest.tools/assert"
 )
 
@@ -30,3 +33,64 @@ func TestWhitelistCheckpoint(t *testing.T) {
 	s.DequeueCheckpointWhitelist()
 	assert.Equal(t, len(s.GetCheckpointWhitelist()), 10, "expected 10 items in whitelist")
 }
+
+// TestIsValidChain checks che IsValidChain function in isolation
+// for different cases by providing a mock fetchHeadersByNumber function
+func TestIsValidChain(t *testing.T) {
+	t.Parallel()
+
+	s := NewMockService(10)
+
+	// case1: no checkpoint whitelist, should consider the chain as valid
+	res, err := s.IsValidChain(nil, nil)
+	assert.Equal(t, res, true, "expected chain to be valid")
+	assert.NilError(t, err, "expected no error")
+
+	// add checkpoint entries and mock fetchHeadersByNumber function
+	s.ProcessCheckpoint(uint64(0), common.Hash{})
+	s.ProcessCheckpoint(uint64(1), common.Hash{})
+	assert.Equal(t, len(s.GetCheckpointWhitelist()), 2, "expected 2 items in whitelist")
+
+	// create a false function, returning absolutely nothing
+	falseFetchHeadersByNumber := func(number uint64, amount int, skip int, reverse bool) ([]*types.Header, []common.Hash, error) {
+		return nil, nil, nil
+	}
+
+	// create a mock function, returning a the required header
+	fetchHeadersByNumber := func(number uint64, amount int, skip int, reverse bool) ([]*types.Header, []common.Hash, error) {
+		hash := common.Hash{}
+		header := types.Header{Number: big.NewInt(0)}
+		switch number {
+		case 0:
+			return []*types.Header{&header}, []common.Hash{hash}, nil
+		case 1:
+			header.Number = big.NewInt(1)
+			return []*types.Header{&header}, []common.Hash{hash}, nil
+		case 2:
+			header.Number = big.NewInt(1) // sending wrong header for misamatch
+			return []*types.Header{&header}, []common.Hash{hash}, nil
+		default:
+			return nil, nil, errors.New("invalid number")
+		}
+	}
+
+	// case2: false fetchHeadersByNumber function provided, should consider the chain as valid
+	res, err = s.IsValidChain(nil, falseFetchHeadersByNumber)
+	assert.Equal(t, res, true, "expected chain to be valid")
+	assert.NilError(t, err, "expected no error")
+
+	// case3: correct fetchHeadersByNumber function provided, should consider the chain as valid
+	res, err = s.IsValidChain(nil, fetchHeadersByNumber)
+	assert.Equal(t, res, true, "expected chain to be valid")
+	assert.NilError(t, err, "expected no error")
+
+	// add one more checkpoint whitelist entry
+	s.ProcessCheckpoint(uint64(2), common.Hash{})
+	assert.Equal(t, len(s.GetCheckpointWhitelist()), 3, "expected 3 items in whitelist")
+
+	// case4: correct fetchHeadersByNumber function provided with wrong header
+	// for block number 2. Should consider the chain as invalid and throw an error
+	res, err = s.IsValidChain(nil, fetchHeadersByNumber)
+	assert.Equal(t, res, false, "expected chain to be invalid")
+	assert.Equal(t, err, ErrCheckpointMismatch, "expected checkpoint mismatch error")
+}