Unit tests for new process_withdrawal logic in electra

Author: prestonvanloon

beacon-chain/core/blocks/withdrawals_test.go

@@ -12,13 +12,15 @@ import (
 	fieldparams "github.com/prysmaticlabs/prysm/v5/config/fieldparams"
 	"github.com/prysmaticlabs/prysm/v5/config/params"
 	consensusblocks "github.com/prysmaticlabs/prysm/v5/consensus-types/blocks"
+	"github.com/prysmaticlabs/prysm/v5/consensus-types/interfaces"
 	"github.com/prysmaticlabs/prysm/v5/consensus-types/primitives"
 	"github.com/prysmaticlabs/prysm/v5/crypto/bls"
 	"github.com/prysmaticlabs/prysm/v5/crypto/bls/common"
 	"github.com/prysmaticlabs/prysm/v5/crypto/hash"
 	"github.com/prysmaticlabs/prysm/v5/encoding/ssz"
 	enginev1 "github.com/prysmaticlabs/prysm/v5/proto/engine/v1"
 	ethpb "github.com/prysmaticlabs/prysm/v5/proto/prysm/v1alpha1"
+	"github.com/prysmaticlabs/prysm/v5/runtime/version"
 	"github.com/prysmaticlabs/prysm/v5/testing/require"
 	"github.com/prysmaticlabs/prysm/v5/time/slots"
 )
@@ -675,6 +677,7 @@ func TestProcessWithdrawals(t *testing.T) {
 		FullWithdrawalIndices           []primitives.ValidatorIndex
 		PendingPartialWithdrawalIndices []primitives.ValidatorIndex
 		Withdrawals                     []*enginev1.Withdrawal
+		PendingPartialWithdrawals       []*ethpb.PendingPartialWithdrawal // Electra
 	}
 	type control struct {
 		NextWithdrawalValidatorIndex primitives.ValidatorIndex
@@ -772,7 +775,7 @@ func TestProcessWithdrawals(t *testing.T) {
 		},
 		{
 			Args: args{
-				Name:                         "Less than max sweep at end",
+				Name:                         "less than max sweep at end",
 				NextWithdrawalIndex:          22,
 				NextWithdrawalValidatorIndex: 4,
 				FullWithdrawalIndices:        []primitives.ValidatorIndex{80, 81, 82, 83},
@@ -789,7 +792,7 @@ func TestProcessWithdrawals(t *testing.T) {
 		},
 		{
 			Args: args{
-				Name:                         "Less than max sweep and beginning",
+				Name:                         "less than max sweep and beginning",
 				NextWithdrawalIndex:          22,
 				NextWithdrawalValidatorIndex: 4,
 				FullWithdrawalIndices:        []primitives.ValidatorIndex{4, 5, 6},
@@ -846,6 +849,36 @@ func TestProcessWithdrawals(t *testing.T) {
 				},
 			},
 		},
+    {
+			Args: args{
+				Name:                            "success many withdrawals with pending partial withdrawals in state",
+				NextWithdrawalIndex:             22,
+				NextWithdrawalValidatorIndex:    88,
+				FullWithdrawalIndices:           []primitives.ValidatorIndex{7, 19, 28},
+				PendingPartialWithdrawalIndices: []primitives.ValidatorIndex{2, 1, 89, 15},
+				Withdrawals: []*enginev1.Withdrawal{
+					PendingPartialWithdrawal(89, 22), PendingPartialWithdrawal(1, 23), PendingPartialWithdrawal(2, 24),
+					fullWithdrawal(7, 25), PendingPartialWithdrawal(15, 26), fullWithdrawal(19, 27),
+					fullWithdrawal(28, 28),
+				},
+				PendingPartialWithdrawals: []*ethpb.PendingPartialWithdrawal{
+				  {
+            Index: 11,
+						Amount: withdrawalAmount(11) - maxEffectiveBalance,
+					},	
+				},
+			},
+			Control: control{
+				NextWithdrawalValidatorIndex: 40,
+				NextWithdrawalIndex:          29,
+				Balances: map[uint64]uint64{
+					7: 0, 19: 0, 28: 0,
+					2: maxEffectiveBalance, 1: maxEffectiveBalance, 89: maxEffectiveBalance,
+					15: maxEffectiveBalance,
+				},
+			},
+		},
+
 		{
 			Args: args{
 				Name:                         "success more than max fully withdrawals",
@@ -1011,65 +1044,97 @@ func TestProcessWithdrawals(t *testing.T) {
 		}
 	}
 
-	prepareValidators := func(st *ethpb.BeaconStateCapella, arguments args) (state.BeaconState, error) {
+	prepareValidators := func(st state.BeaconState, arguments args) error {
 		validators := make([]*ethpb.Validator, numValidators)
-		st.Balances = make([]uint64, numValidators)
+		if err := st.SetBalances(make([]uint64, numValidators)); err != nil {
+			return err
+		}
 		for i := range validators {
 			v := &ethpb.Validator{}
 			v.EffectiveBalance = maxEffectiveBalance
 			v.WithdrawableEpoch = epochInFuture
 			v.WithdrawalCredentials = make([]byte, 32)
 			v.WithdrawalCredentials[31] = byte(i)
-			st.Balances[i] = v.EffectiveBalance - uint64(rand.Intn(1000))
+			if err := st.UpdateBalancesAtIndex(primitives.ValidatorIndex(i), v.EffectiveBalance-uint64(rand.Intn(1000))); err != nil {
+				return err
+			}
 			validators[i] = v
 		}
 		for _, idx := range arguments.FullWithdrawalIndices {
 			if idx != notWithdrawableIndex {
 				validators[idx].WithdrawableEpoch = epochInPast
 			}
-			st.Balances[idx] = withdrawalAmount(idx)
+			if err := st.UpdateBalancesAtIndex(idx, withdrawalAmount(idx)); err != nil {
+				return err
+			}
 			validators[idx].WithdrawalCredentials[0] = params.BeaconConfig().ETH1AddressWithdrawalPrefixByte
 		}
 		for _, idx := range arguments.PendingPartialWithdrawalIndices {
 			validators[idx].WithdrawalCredentials[0] = params.BeaconConfig().ETH1AddressWithdrawalPrefixByte
-			st.Balances[idx] = withdrawalAmount(idx)
+			if err := st.UpdateBalancesAtIndex(idx, withdrawalAmount(idx)); err != nil {
+				return err
+			}
 		}
-		st.Validators = validators
-		return state_native.InitializeFromProtoCapella(st)
+		return st.SetValidators(validators)
 	}
 
 	for _, test := range tests {
 		t.Run(test.Args.Name, func(t *testing.T) {
-			saved := params.BeaconConfig().MaxValidatorsPerWithdrawalsSweep
-			params.BeaconConfig().MaxValidatorsPerWithdrawalsSweep = maxSweep
-			if test.Args.Withdrawals == nil {
-				test.Args.Withdrawals = make([]*enginev1.Withdrawal, 0)
-			}
-			if test.Args.FullWithdrawalIndices == nil {
-				test.Args.FullWithdrawalIndices = make([]primitives.ValidatorIndex, 0)
-			}
-			if test.Args.PendingPartialWithdrawalIndices == nil {
-				test.Args.PendingPartialWithdrawalIndices = make([]primitives.ValidatorIndex, 0)
-			}
-			slot, err := slots.EpochStart(currentEpoch)
-			require.NoError(t, err)
-			spb := &ethpb.BeaconStateCapella{
-				Slot:                         slot,
-				NextWithdrawalValidatorIndex: test.Args.NextWithdrawalValidatorIndex,
-				NextWithdrawalIndex:          test.Args.NextWithdrawalIndex,
-			}
-			st, err := prepareValidators(spb, test.Args)
-			require.NoError(t, err)
-			p, err := consensusblocks.WrappedExecutionPayloadCapella(&enginev1.ExecutionPayloadCapella{Withdrawals: test.Args.Withdrawals})
-			require.NoError(t, err)
-			post, err := blocks.ProcessWithdrawals(st, p)
-			if test.Control.ExpectedError {
-				require.NotNil(t, err)
-			} else {
-				require.NoError(t, err)
-				checkPostState(t, test.Control, post)
+			for _, fork := range []int{version.Capella, version.Electra} {
+				t.Run(version.String(fork), func(t *testing.T) {
+					saved := params.BeaconConfig().MaxValidatorsPerWithdrawalsSweep
+					params.BeaconConfig().MaxValidatorsPerWithdrawalsSweep = maxSweep
+					if test.Args.Withdrawals == nil {
+						test.Args.Withdrawals = make([]*enginev1.Withdrawal, 0)
+					}
+					if test.Args.FullWithdrawalIndices == nil {
+						test.Args.FullWithdrawalIndices = make([]primitives.ValidatorIndex, 0)
+					}
+					if test.Args.PendingPartialWithdrawalIndices == nil {
+						test.Args.PendingPartialWithdrawalIndices = make([]primitives.ValidatorIndex, 0)
+					}
+					slot, err := slots.EpochStart(currentEpoch)
+					require.NoError(t, err)
+					var st state.BeaconState
+					var p interfaces.ExecutionData
+					switch fork {
+					case version.Capella:
+						spb := &ethpb.BeaconStateCapella{
+							Slot:                         slot,
+							NextWithdrawalValidatorIndex: test.Args.NextWithdrawalValidatorIndex,
+							NextWithdrawalIndex:          test.Args.NextWithdrawalIndex,
+						}
+						st, err = state_native.InitializeFromProtoUnsafeCapella(spb)
+						require.NoError(t, err)
+						p, err = consensusblocks.WrappedExecutionPayloadCapella(&enginev1.ExecutionPayloadCapella{Withdrawals: test.Args.Withdrawals})
+						require.NoError(t, err)
+					case version.Electra:
+						spb := &ethpb.BeaconStateElectra{
+							Slot:                         slot,
+							NextWithdrawalValidatorIndex: test.Args.NextWithdrawalValidatorIndex,
+							NextWithdrawalIndex:          test.Args.NextWithdrawalIndex,
+							PendingPartialWithdrawals:    test.Args.PendingPartialWithdrawals,
+						}
+						st, err = state_native.InitializeFromProtoUnsafeElectra(spb)
+						require.NoError(t, err)
+						p, err = consensusblocks.WrappedExecutionPayloadElectra(&enginev1.ExecutionPayloadElectra{Withdrawals: test.Args.Withdrawals})
+						require.NoError(t, err)
+					default:
+						t.Fatalf("Add a beacon state setup for version %s", version.String(fork))
+					}
+					err = prepareValidators(st, test.Args)
+					require.NoError(t, err)
+					post, err := blocks.ProcessWithdrawals(st, p)
+					if test.Control.ExpectedError {
+						require.NotNil(t, err)
+					} else {
+						require.NoError(t, err)
+						checkPostState(t, test.Control, post)
+					}
+					params.BeaconConfig().MaxValidatorsPerWithdrawalsSweep = saved
+
+				})
 			}
-			params.BeaconConfig().MaxValidatorsPerWithdrawalsSweep = saved
 		})
 	}
 }