Spectest: refactor forkchoice helper (#10616)

* Seperate forkchoice builder for lower cognitive score and better reusablity

* gaz

* deepsource resolutions

Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>

testing/spectest/shared/common/forkchoice/BUILD.bazel

@@ -1,9 +1,10 @@
-load("@prysm//tools/go:def.bzl", "go_library")
+load("@prysm//tools/go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "go_default_library",
     testonly = True,
     srcs = [
+        "builder.go",
         "runner.go",
         "service.go",
         "type.go",
@@ -43,3 +44,15 @@ go_library(
         "@com_github_golang_snappy//:go_default_library",
     ],
 )
+
+go_test(
+    name = "go_default_test",
+    srcs = ["builder_test.go"],
+    embed = [":go_default_library"],
+    deps = [
+        "//consensus-types/wrapper:go_default_library",
+        "//proto/prysm/v1alpha1:go_default_library",
+        "//testing/require:go_default_library",
+        "//testing/util:go_default_library",
+    ],
+)

testing/spectest/shared/common/forkchoice/builder.go

@@ -0,0 +1,127 @@
+package forkchoice
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/prysmaticlabs/prysm/beacon-chain/blockchain"
+	forkchoicetypes "github.com/prysmaticlabs/prysm/beacon-chain/forkchoice/types"
+	"github.com/prysmaticlabs/prysm/beacon-chain/state"
+	"github.com/prysmaticlabs/prysm/config/params"
+	"github.com/prysmaticlabs/prysm/consensus-types/interfaces"
+	types "github.com/prysmaticlabs/prysm/consensus-types/primitives"
+	"github.com/prysmaticlabs/prysm/encoding/bytesutil"
+	ethpb "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1"
+	"github.com/prysmaticlabs/prysm/testing/require"
+	"github.com/prysmaticlabs/prysm/time/slots"
+)
+
+type Builder struct {
+	service  *blockchain.Service
+	lastTick int64
+	execMock *engineMock
+}
+
+func NewBuilder(t testing.TB, initialState state.BeaconState, initialBlock interfaces.SignedBeaconBlock) *Builder {
+	execMock := &engineMock{
+		powBlocks: make(map[[32]byte]*ethpb.PowBlock),
+	}
+	service := startChainService(t, initialState, initialBlock, execMock)
+	return &Builder{
+		service:  service,
+		execMock: execMock,
+	}
+}
+
+// Tick resets the genesis time to now()-tick and adjusts the slot to the appropriate value.
+func (bb *Builder) Tick(t testing.TB, tick int64) {
+	bb.service.SetGenesisTime(time.Unix(time.Now().Unix()-tick, 0))
+	if tick > bb.lastTick {
+		slot := uint64(tick) / params.BeaconConfig().SecondsPerSlot
+		require.NoError(t, bb.service.NewSlot(context.TODO(), types.Slot(slot)))
+		bb.lastTick = tick
+	}
+}
+
+// block provides the block to forkchoice proposer boost and returns the block root.
+func (bb *Builder) block(t testing.TB, b interfaces.SignedBeaconBlock) [32]byte {
+	ctx := context.TODO()
+	r, err := b.Block().HashTreeRoot()
+	require.NoError(t, err)
+	slotsSinceGenesis := slots.SinceGenesis(bb.service.GenesisTime())
+	args := &forkchoicetypes.ProposerBoostRootArgs{
+		BlockRoot:       r,
+		BlockSlot:       b.Block().Slot(),
+		CurrentSlot:     slotsSinceGenesis,
+		SecondsIntoSlot: uint64(bb.lastTick) % params.BeaconConfig().SecondsPerSlot,
+	}
+	require.NoError(t, bb.service.ForkChoicer().BoostProposerRoot(ctx, args))
+	return r
+}
+
+// InvalidBlock receives the invalid block and notifies forkchoice.
+func (bb *Builder) InvalidBlock(t testing.TB, b interfaces.SignedBeaconBlock) {
+	r := bb.block(t, b)
+	require.Equal(t, true, bb.service.ReceiveBlock(context.TODO(), b, r) != nil)
+}
+
+// ValidBlock receives the valid block and notifies forkchoice.
+func (bb *Builder) ValidBlock(t testing.TB, b interfaces.SignedBeaconBlock) {
+	r := bb.block(t, b)
+	require.NoError(t, bb.service.ReceiveBlock(context.TODO(), b, r))
+}
+
+// PoWBlock receives the block and notifies a mocked execution engine.
+func (bb *Builder) PoWBlock(t testing.TB, pb *ethpb.PowBlock) {
+	bb.execMock.powBlocks[bytesutil.ToBytes32(pb.BlockHash)] = pb
+}
+
+// Attestation receives the attestation and updates forkchoice.
+func (bb *Builder) Attestation(t testing.TB, a *ethpb.Attestation) {
+	require.NoError(t, bb.service.OnAttestation(context.TODO(), a))
+}
+
+// Check evaluates the fork choice results and compares them to the expected values.
+func (bb *Builder) Check(t testing.TB, c *Check) {
+	if c == nil {
+		return
+	}
+	ctx := context.TODO()
+
+	require.NoError(t, bb.service.UpdateAndSaveHeadWithBalances(ctx))
+	if c.Head != nil {
+		r, err := bb.service.HeadRoot(ctx)
+		require.NoError(t, err)
+		require.DeepEqual(t, common.FromHex(c.Head.Root), r)
+		require.Equal(t, types.Slot(c.Head.Slot), bb.service.HeadSlot())
+	}
+	if c.JustifiedCheckPoint != nil {
+		cp := &ethpb.Checkpoint{
+			Epoch: types.Epoch(c.JustifiedCheckPoint.Epoch),
+			Root:  common.FromHex(c.JustifiedCheckPoint.Root),
+		}
+		require.DeepEqual(t, cp, bb.service.CurrentJustifiedCheckpt())
+	}
+	if c.BestJustifiedCheckPoint != nil {
+		cp := &ethpb.Checkpoint{
+			Epoch: types.Epoch(c.BestJustifiedCheckPoint.Epoch),
+			Root:  common.FromHex(c.BestJustifiedCheckPoint.Root),
+		}
+		require.DeepEqual(t, cp, bb.service.BestJustifiedCheckpt())
+	}
+	if c.FinalizedCheckPoint != nil {
+		cp := &ethpb.Checkpoint{
+			Epoch: types.Epoch(c.FinalizedCheckPoint.Epoch),
+			Root:  common.FromHex(c.FinalizedCheckPoint.Root),
+		}
+		require.DeepSSZEqual(t, cp, bb.service.FinalizedCheckpt())
+	}
+	if c.ProposerBoostRoot != nil {
+		want := fmt.Sprintf("%#x", common.FromHex(*c.ProposerBoostRoot))
+		got := fmt.Sprintf("%#x", bb.service.ForkChoiceStore().ProposerBoost())
+		require.DeepEqual(t, want, got)
+	}
+}

testing/spectest/shared/common/forkchoice/builder_test.go

@@ -0,0 +1,41 @@
+package forkchoice
+
+import (
+	"testing"
+
+	"github.com/prysmaticlabs/prysm/consensus-types/wrapper"
+	ethpb "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1"
+	"github.com/prysmaticlabs/prysm/testing/require"
+	"github.com/prysmaticlabs/prysm/testing/util"
+)
+
+func TestBuilderTick(t *testing.T) {
+	st, err := util.NewBeaconState()
+	require.NoError(t, err)
+	blk, err := wrapper.WrappedSignedBeaconBlock(util.NewBeaconBlock())
+	require.NoError(t, err)
+	builder := NewBuilder(t, st, blk)
+	builder.Tick(t, 10)
+
+	require.Equal(t, int64(10), builder.lastTick)
+}
+
+func TestBuilderInvalidBlock(t *testing.T) {
+	st, err := util.NewBeaconState()
+	require.NoError(t, err)
+	blk, err := wrapper.WrappedSignedBeaconBlock(util.NewBeaconBlock())
+	require.NoError(t, err)
+	builder := NewBuilder(t, st, blk)
+	builder.InvalidBlock(t, blk)
+}
+
+func TestPoWBlock(t *testing.T) {
+	st, err := util.NewBeaconState()
+	require.NoError(t, err)
+	blk, err := wrapper.WrappedSignedBeaconBlock(util.NewBeaconBlock())
+	require.NoError(t, err)
+	builder := NewBuilder(t, st, blk)
+	builder.PoWBlock(t, &ethpb.PowBlock{BlockHash: []byte{1, 2, 3}})
+
+	require.Equal(t, 1, len(builder.execMock.powBlocks))
+}

testing/spectest/shared/common/forkchoice/runner.go

@@ -1,32 +1,23 @@
 package forkchoice
 
 import (
-	"context"
 	"fmt"
-	"math/big"
 	"path"
 	"testing"
-	"time"
 
-	"github.com/ethereum/go-ethereum/common"
 	"github.com/golang/snappy"
-	forkchoicetypes "github.com/prysmaticlabs/prysm/beacon-chain/forkchoice/types"
 	"github.com/prysmaticlabs/prysm/beacon-chain/state"
 	v1 "github.com/prysmaticlabs/prysm/beacon-chain/state/v1"
 	v2 "github.com/prysmaticlabs/prysm/beacon-chain/state/v2"
 	v3 "github.com/prysmaticlabs/prysm/beacon-chain/state/v3"
 	fieldparams "github.com/prysmaticlabs/prysm/config/fieldparams"
-	"github.com/prysmaticlabs/prysm/config/params"
 	"github.com/prysmaticlabs/prysm/consensus-types/interfaces"
-	types "github.com/prysmaticlabs/prysm/consensus-types/primitives"
 	"github.com/prysmaticlabs/prysm/consensus-types/wrapper"
-	"github.com/prysmaticlabs/prysm/encoding/bytesutil"
 	ethpb "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1"
 	"github.com/prysmaticlabs/prysm/runtime/version"
 	"github.com/prysmaticlabs/prysm/testing/require"
 	"github.com/prysmaticlabs/prysm/testing/spectest/utils"
 	"github.com/prysmaticlabs/prysm/testing/util"
-	"github.com/prysmaticlabs/prysm/time/slots"
 )
 
 // Run executes "forkchoice" test.
@@ -40,7 +31,6 @@ func Run(t *testing.T, config string, fork int) {
 
 		for _, folder := range testFolders {
 			t.Run(folder.Name(), func(t *testing.T) {
-				ctx := context.Background()
 				preStepsFile, err := util.BazelFileBytes(testsFolderPath, folder.Name(), "steps.yaml")
 				require.NoError(t, err)
 				var steps []Step
@@ -72,20 +62,11 @@ func Run(t *testing.T, config string, fork int) {
 					t.Fatalf("unknown fork version: %v", fork)
 				}
 
-				execMock := &engineMock{
-					powBlocks: make(map[[32]byte]*ethpb.PowBlock),
-				}
-				service := startChainService(t, beaconState, beaconBlock, execMock)
-				var lastTick int64
+				builder := NewBuilder(t, beaconState, beaconBlock)
+
 				for _, step := range steps {
 					if step.Tick != nil {
-						newTick := int64(*step.Tick)
-						service.SetGenesisTime(time.Unix(time.Now().Unix()-newTick, 0))
-						if newTick > lastTick {
-							slot := uint64(newTick) / params.BeaconConfig().SecondsPerSlot
-							require.NoError(t, service.NewSlot(ctx, types.Slot(slot)))
-							lastTick = newTick
-						}
+						builder.Tick(t, int64(*step.Tick))
 					}
 					if step.Block != nil {
 						blockFile, err := util.BazelFileBytes(testsFolderPath, folder.Name(), fmt.Sprint(*step.Block, ".ssz_snappy"))
@@ -103,20 +84,10 @@ func Run(t *testing.T, config string, fork int) {
 						default:
 							t.Fatalf("unknown fork version: %v", fork)
 						}
-						r, err := beaconBlock.Block().HashTreeRoot()
-						require.NoError(t, err)
-						slotsSinceGenesis := slots.SinceGenesis(service.GenesisTime())
-						args := &forkchoicetypes.ProposerBoostRootArgs{
-							BlockRoot:       r,
-							BlockSlot:       beaconBlock.Block().Slot(),
-							CurrentSlot:     slotsSinceGenesis,
-							SecondsIntoSlot: uint64(lastTick) % params.BeaconConfig().SecondsPerSlot,
-						}
-						require.NoError(t, service.ForkChoicer().BoostProposerRoot(ctx, args))
 						if step.Valid != nil && !*step.Valid {
-							require.Equal(t, true, service.ReceiveBlock(ctx, beaconBlock, r) != nil)
+							builder.InvalidBlock(t, beaconBlock)
 						} else {
-							require.NoError(t, service.ReceiveBlock(ctx, beaconBlock, r))
+							builder.ValidBlock(t, beaconBlock)
 						}
 					}
 					if step.Attestation != nil {
@@ -126,7 +97,7 @@ func Run(t *testing.T, config string, fork int) {
 						require.NoError(t, err)
 						att := &ethpb.Attestation{}
 						require.NoError(t, att.UnmarshalSSZ(attSSZ), "Failed to unmarshal")
-						require.NoError(t, service.OnAttestation(ctx, att))
+						builder.Attestation(t, att)
 					}
 					if step.PowBlock != nil {
 						powBlockFile, err := util.BazelFileBytes(testsFolderPath, folder.Name(), fmt.Sprint(*step.PowBlock, ".ssz_snappy"))
@@ -135,47 +106,9 @@ func Run(t *testing.T, config string, fork int) {
 						require.NoError(t, err)
 						pb := &ethpb.PowBlock{}
 						require.NoError(t, pb.UnmarshalSSZ(p), "Failed to unmarshal")
-						execMock.powBlocks[bytesutil.ToBytes32(pb.BlockHash)] = pb
-						tdInBigEndian := bytesutil.ReverseByteOrder(pb.TotalDifficulty)
-						tdBigint := new(big.Int)
-						tdBigint.SetBytes(tdInBigEndian)
-					}
-					if step.Check != nil {
-						require.NoError(t, service.UpdateAndSaveHeadWithBalances(ctx))
-						c := step.Check
-						if c.Head != nil {
-							r, err := service.HeadRoot(ctx)
-							require.NoError(t, err)
-							require.DeepEqual(t, common.FromHex(c.Head.Root), r)
-							require.Equal(t, types.Slot(c.Head.Slot), service.HeadSlot())
-						}
-						if c.JustifiedCheckPoint != nil {
-							cp := &ethpb.Checkpoint{
-								Epoch: types.Epoch(c.JustifiedCheckPoint.Epoch),
-								Root:  common.FromHex(c.JustifiedCheckPoint.Root),
-							}
-							require.DeepEqual(t, cp, service.CurrentJustifiedCheckpt())
-						}
-						if c.BestJustifiedCheckPoint != nil {
-							cp := &ethpb.Checkpoint{
-								Epoch: types.Epoch(c.BestJustifiedCheckPoint.Epoch),
-								Root:  common.FromHex(c.BestJustifiedCheckPoint.Root),
-							}
-							require.DeepEqual(t, cp, service.BestJustifiedCheckpt())
-						}
-						if c.FinalizedCheckPoint != nil {
-							cp := &ethpb.Checkpoint{
-								Epoch: types.Epoch(c.FinalizedCheckPoint.Epoch),
-								Root:  common.FromHex(c.FinalizedCheckPoint.Root),
-							}
-							require.DeepSSZEqual(t, cp, service.FinalizedCheckpt())
-						}
-						if c.ProposerBoostRoot != nil {
-							want := fmt.Sprintf("%#x", common.FromHex(*c.ProposerBoostRoot))
-							got := fmt.Sprintf("%#x", service.ForkChoiceStore().ProposerBoost())
-							require.DeepEqual(t, want, got)
-						}
+						builder.PoWBlock(t, pb)
 					}
+					builder.Check(t, step.Check)
 				}
 			})
 		}

testing/spectest/shared/common/forkchoice/service.go

@@ -25,7 +25,7 @@ import (
 	"github.com/prysmaticlabs/prysm/testing/require"
 )
 
-func startChainService(t *testing.T, st state.BeaconState, block interfaces.SignedBeaconBlock, engineMock *engineMock) *blockchain.Service {
+func startChainService(t testing.TB, st state.BeaconState, block interfaces.SignedBeaconBlock, engineMock *engineMock) *blockchain.Service {
 	db := testDB.SetupDB(t)
 	ctx := context.Background()
 	require.NoError(t, db.SaveBlock(ctx, block))