From ae696b64bafea8971a8ea118b433c5efb58f89c2 Mon Sep 17 00:00:00 2001
From: Preston Van Loon <preston@pvl.dev>
Date: Tue, 2 Jul 2024 14:23:20 -0500
Subject: [PATCH] Add unit test for VerifyExitAndSignature EIP-7251

---
 beacon-chain/core/blocks/exit.go      |  2 +-
 beacon-chain/core/blocks/exit_test.go | 53 ++++++++++++++++++++++++++-
 2 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/beacon-chain/core/blocks/exit.go b/beacon-chain/core/blocks/exit.go
index 4b789613052a..be0bfd3bb7be 100644
--- a/beacon-chain/core/blocks/exit.go
+++ b/beacon-chain/core/blocks/exit.go
@@ -190,7 +190,7 @@ func verifyExitConditions(st state.ReadOnlyBeaconState, validator state.ReadOnly
 			params.BeaconConfig().ShardCommitteePeriod,
 			validator.ActivationEpoch()+params.BeaconConfig().ShardCommitteePeriod,
 		)
-	} 
+	}
 
 	if st.Version() >= version.Electra {
 		// Only exit validator if it has no pending withdrawals in the queue.
diff --git a/beacon-chain/core/blocks/exit_test.go b/beacon-chain/core/blocks/exit_test.go
index d40f35f7a2be..bda5f7ef81b4 100644
--- a/beacon-chain/core/blocks/exit_test.go
+++ b/beacon-chain/core/blocks/exit_test.go
@@ -135,6 +135,11 @@ func TestProcessVoluntaryExits_AppliesCorrectStatus(t *testing.T) {
 }
 
 func TestVerifyExitAndSignature(t *testing.T) {
+	// Remove after electra fork epoch is defined.
+	cfg := params.BeaconConfig()
+	cfg.ElectraForkEpoch = cfg.DenebForkEpoch * 2
+	params.SetActiveTestCleanup(t, cfg)
+	// End remove section.
 	denebSlot, err := slots.EpochStart(params.BeaconConfig().DenebForkEpoch)
 	require.NoError(t, err)
 	tests := []struct {
@@ -167,7 +172,7 @@ func TestVerifyExitAndSignature(t *testing.T) {
 			wantErr: "nil exit",
 		},
 		{
-			name: "Happy Path",
+			name: "Happy Path phase0",
 			setup: func() (*ethpb.Validator, *ethpb.SignedVoluntaryExit, state.ReadOnlyBeaconState, error) {
 				fork := &ethpb.Fork{
 					PreviousVersion: params.BeaconConfig().GenesisForkVersion,
@@ -275,6 +280,52 @@ func TestVerifyExitAndSignature(t *testing.T) {
 				return validator, signedExit, bs, nil
 			},
 		},
+		{
+			name: "EIP-7251 - pending balance to withdraw must be zero",
+			setup: func() (*ethpb.Validator, *ethpb.SignedVoluntaryExit, state.ReadOnlyBeaconState, error) {
+				fork := &ethpb.Fork{
+					PreviousVersion: params.BeaconConfig().DenebForkVersion,
+					CurrentVersion:  params.BeaconConfig().ElectraForkVersion,
+					Epoch:           params.BeaconConfig().ElectraForkEpoch,
+				}
+				signedExit := &ethpb.SignedVoluntaryExit{
+					Exit: &ethpb.VoluntaryExit{
+						Epoch:          params.BeaconConfig().DenebForkEpoch + 1,
+						ValidatorIndex: 0,
+					},
+				}
+				electraSlot, err := slots.EpochStart(params.BeaconConfig().ElectraForkEpoch)
+				require.NoError(t, err)
+				bs, keys := util.DeterministicGenesisState(t, 1)
+				bs, err = state_native.InitializeFromProtoUnsafeElectra(&ethpb.BeaconStateElectra{
+					GenesisValidatorsRoot: bs.GenesisValidatorsRoot(),
+					Fork:                  fork,
+					Slot:                  electraSlot,
+					Validators:            bs.Validators(),
+				})
+				if err != nil {
+					return nil, nil, nil, err
+				}
+				validator := bs.Validators()[0]
+				validator.ActivationEpoch = 1
+				err = bs.UpdateValidatorAtIndex(0, validator)
+				require.NoError(t, err)
+				sb, err := signing.ComputeDomainAndSign(bs, signedExit.Exit.Epoch, signedExit.Exit, params.BeaconConfig().DomainVoluntaryExit, keys[0])
+				require.NoError(t, err)
+				sig, err := bls.SignatureFromBytes(sb)
+				require.NoError(t, err)
+				signedExit.Signature = sig.Marshal()
+
+				// Give validator a pending balance to withdraw.
+				require.NoError(t, bs.AppendPendingPartialWithdrawal(&ethpb.PendingPartialWithdrawal{
+					Index:  0,
+					Amount: 500,
+				}))
+
+				return validator, signedExit, bs, nil
+			},
+			wantErr: "validator 0 must have no pending balance to withdraw",
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {