Refactor for DRY state getters_validator_test and getters_test

beacon-chain/state/state-native/v1/BUILD.bazel

@@ -97,6 +97,7 @@ go_test(
     deps = [
         "//beacon-chain/state:go_default_library",
         "//beacon-chain/state/stateutil:go_default_library",
+        "//beacon-chain/state/testing:go_default_library",
         "//beacon-chain/state/types:go_default_library",
         "//config/features:go_default_library",
         "//config/fieldparams:go_default_library",

beacon-chain/state/state-native/v1/getters_test.go

@@ -1,155 +1,39 @@
 package v1
 
 import (
-	"sync"
 	"testing"
 
-	types "github.com/prysmaticlabs/eth2-types"
-	fieldparams "github.com/prysmaticlabs/prysm/config/fieldparams"
+	"github.com/prysmaticlabs/prysm/beacon-chain/state"
+	testtmpl "github.com/prysmaticlabs/prysm/beacon-chain/state/testing"
 	ethpb "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1"
-	"github.com/prysmaticlabs/prysm/testing/assert"
-	"github.com/prysmaticlabs/prysm/testing/require"
 )
 
 func TestBeaconState_SlotDataRace(t *testing.T) {
-	headState, err := InitializeFromProto(&ethpb.BeaconState{Slot: 1})
-	require.NoError(t, err)
-
-	wg := sync.WaitGroup{}
-	wg.Add(2)
-	go func() {
-		require.NoError(t, headState.SetSlot(0))
-		wg.Done()
-	}()
-	go func() {
-		headState.Slot()
-		wg.Done()
-	}()
-
-	wg.Wait()
+	testtmpl.VerifyBeaconState_SlotDataRace(t, func() (state.BeaconState, error) {
+		return InitializeFromProto(&ethpb.BeaconState{Slot: 1})
+	})
 }
 
 func TestBeaconState_MatchCurrentJustifiedCheckpt(t *testing.T) {
-	c1 := &ethpb.Checkpoint{Epoch: 1}
-	c2 := &ethpb.Checkpoint{Epoch: 2}
-	beaconState, err := InitializeFromProto(&ethpb.BeaconState{CurrentJustifiedCheckpoint: c1})
-	require.NoError(t, err)
-	require.Equal(t, true, beaconState.MatchCurrentJustifiedCheckpoint(c1))
-	require.Equal(t, false, beaconState.MatchCurrentJustifiedCheckpoint(c2))
-	require.Equal(t, false, beaconState.MatchPreviousJustifiedCheckpoint(c1))
-	require.Equal(t, false, beaconState.MatchPreviousJustifiedCheckpoint(c2))
+	testtmpl.VerifyBeaconState_MatchCurrentJustifiedCheckptNative(
+		t,
+		func(cp *ethpb.Checkpoint) (state.BeaconState, error) {
+			return InitializeFromProto(&ethpb.BeaconState{CurrentJustifiedCheckpoint: cp})
+		},
+	)
 }
 
 func TestBeaconState_MatchPreviousJustifiedCheckpt(t *testing.T) {
-	c1 := &ethpb.Checkpoint{Epoch: 1}
-	c2 := &ethpb.Checkpoint{Epoch: 2}
-	beaconState, err := InitializeFromProto(&ethpb.BeaconState{PreviousJustifiedCheckpoint: c1})
-	require.NoError(t, err)
-	require.NoError(t, err)
-	require.Equal(t, false, beaconState.MatchCurrentJustifiedCheckpoint(c1))
-	require.Equal(t, false, beaconState.MatchCurrentJustifiedCheckpoint(c2))
-	require.Equal(t, true, beaconState.MatchPreviousJustifiedCheckpoint(c1))
-	require.Equal(t, false, beaconState.MatchPreviousJustifiedCheckpoint(c2))
+	testtmpl.VerifyBeaconState_MatchPreviousJustifiedCheckptNative(
+		t,
+		func(cp *ethpb.Checkpoint) (state.BeaconState, error) {
+			return InitializeFromProto(&ethpb.BeaconState{PreviousJustifiedCheckpoint: cp})
+		},
+	)
 }
 
 func TestBeaconState_ValidatorByPubkey(t *testing.T) {
-	keyCreator := func(input []byte) [fieldparams.BLSPubkeyLength]byte {
-		nKey := [fieldparams.BLSPubkeyLength]byte{}
-		copy(nKey[:1], input)
-		return nKey
-	}
-
-	tests := []struct {
-		name            string
-		modifyFunc      func(b *BeaconState, k [fieldparams.BLSPubkeyLength]byte)
-		exists          bool
-		expectedIdx     types.ValidatorIndex
-		largestIdxInSet types.ValidatorIndex
-	}{
-		{
-			name: "retrieve validator",
-			modifyFunc: func(b *BeaconState, key [fieldparams.BLSPubkeyLength]byte) {
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key[:]}))
-			},
-			exists:      true,
-			expectedIdx: 0,
-		},
-		{
-			name: "retrieve validator with multiple validators from the start",
-			modifyFunc: func(b *BeaconState, key [fieldparams.BLSPubkeyLength]byte) {
-				key1 := keyCreator([]byte{'C'})
-				key2 := keyCreator([]byte{'D'})
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key[:]}))
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key1[:]}))
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key2[:]}))
-			},
-			exists:      true,
-			expectedIdx: 0,
-		},
-		{
-			name: "retrieve validator with multiple validators",
-			modifyFunc: func(b *BeaconState, key [fieldparams.BLSPubkeyLength]byte) {
-				key1 := keyCreator([]byte{'C'})
-				key2 := keyCreator([]byte{'D'})
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key1[:]}))
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key2[:]}))
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key[:]}))
-			},
-			exists:      true,
-			expectedIdx: 2,
-		},
-		{
-			name: "retrieve validator with multiple validators from the start with shared state",
-			modifyFunc: func(b *BeaconState, key [fieldparams.BLSPubkeyLength]byte) {
-				key1 := keyCreator([]byte{'C'})
-				key2 := keyCreator([]byte{'D'})
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key[:]}))
-				_ = b.Copy()
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key1[:]}))
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key2[:]}))
-			},
-			exists:      true,
-			expectedIdx: 0,
-		},
-		{
-			name: "retrieve validator with multiple validators with shared state",
-			modifyFunc: func(b *BeaconState, key [fieldparams.BLSPubkeyLength]byte) {
-				key1 := keyCreator([]byte{'C'})
-				key2 := keyCreator([]byte{'D'})
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key1[:]}))
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key2[:]}))
-				n := b.Copy()
-				// Append to another state
-				assert.NoError(t, n.AppendValidator(&ethpb.Validator{PublicKey: key[:]}))
-
-			},
-			exists:      false,
-			expectedIdx: 0,
-		},
-		{
-			name: "retrieve validator with multiple validators with shared state at boundary",
-			modifyFunc: func(b *BeaconState, key [fieldparams.BLSPubkeyLength]byte) {
-				key1 := keyCreator([]byte{'C'})
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key1[:]}))
-				n := b.Copy()
-				// Append to another state
-				assert.NoError(t, n.AppendValidator(&ethpb.Validator{PublicKey: key[:]}))
-
-			},
-			exists:      false,
-			expectedIdx: 0,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			s, err := InitializeFromProto(&ethpb.BeaconState{})
-			require.NoError(t, err)
-			nKey := keyCreator([]byte{'A'})
-			tt.modifyFunc(s, nKey)
-			idx, ok := s.ValidatorIndexByPubkey(nKey)
-			assert.Equal(t, tt.exists, ok)
-			assert.Equal(t, tt.expectedIdx, idx)
-		})
-	}
+	testtmpl.VerifyBeaconState_ValidatorByPubkey(t, func() (state.BeaconState, error) {
+		return InitializeFromProto(&ethpb.BeaconState{})
+	})
 }

beacon-chain/state/state-native/v1/getters_validator_test.go

@@ -5,17 +5,14 @@ import (
 
 	"github.com/prysmaticlabs/prysm/beacon-chain/state"
 	v1 "github.com/prysmaticlabs/prysm/beacon-chain/state/state-native/v1"
+	testtmpl "github.com/prysmaticlabs/prysm/beacon-chain/state/testing"
 	ethpb "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1"
-	"github.com/prysmaticlabs/prysm/testing/assert"
-	"github.com/prysmaticlabs/prysm/testing/require"
 )
 
 func TestBeaconState_ValidatorAtIndexReadOnly_HandlesNilSlice(t *testing.T) {
-	st, err := v1.InitializeFromProtoUnsafe(&ethpb.BeaconState{
-		Validators: nil,
+	testtmpl.VerifyBeaconState_ValidatorAtIndexReadOnly_HandlesNilSlice(t, func() (state.BeaconState, error) {
+		return v1.InitializeFromProtoUnsafe(&ethpb.BeaconState{
+			Validators: nil,
+		})
 	})
-	require.NoError(t, err)
-
-	_, err = st.ValidatorAtIndexReadOnly(0)
-	assert.Equal(t, state.ErrNilValidatorsInState, err)
 }

beacon-chain/state/state-native/v2/BUILD.bazel

@@ -84,8 +84,8 @@ go_test(
     embed = [":go_default_library"],
     deps = [
         "//beacon-chain/state:go_default_library",
-        "//beacon-chain/state/state-native/v1:go_default_library",
         "//beacon-chain/state/stateutil:go_default_library",
+        "//beacon-chain/state/testing:go_default_library",
         "//beacon-chain/state/v2:go_default_library",
         "//config/features:go_default_library",
         "//config/fieldparams:go_default_library",

beacon-chain/state/state-native/v2/getters_test.go

@@ -1,132 +1,39 @@
 package v2
 
 import (
-	"sync"
 	"testing"
 
-	types "github.com/prysmaticlabs/eth2-types"
-	fieldparams "github.com/prysmaticlabs/prysm/config/fieldparams"
+	"github.com/prysmaticlabs/prysm/beacon-chain/state"
+	testtmpl "github.com/prysmaticlabs/prysm/beacon-chain/state/testing"
 	ethpb "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1"
-	"github.com/prysmaticlabs/prysm/testing/assert"
-	"github.com/prysmaticlabs/prysm/testing/require"
 )
 
 func TestBeaconState_SlotDataRace(t *testing.T) {
-	headState, err := InitializeFromProto(&ethpb.BeaconStateAltair{Slot: 1})
-	require.NoError(t, err)
-
-	wg := sync.WaitGroup{}
-	wg.Add(2)
-	go func() {
-		require.NoError(t, headState.SetSlot(0))
-		wg.Done()
-	}()
-	go func() {
-		headState.Slot()
-		wg.Done()
-	}()
-
-	wg.Wait()
+	testtmpl.VerifyBeaconState_SlotDataRace(t, func() (state.BeaconState, error) {
+		return InitializeFromProto(&ethpb.BeaconStateAltair{Slot: 1})
+	})
 }
 
-func TestBeaconState_ValidatorByPubkey(t *testing.T) {
-	keyCreator := func(input []byte) [fieldparams.BLSPubkeyLength]byte {
-		nKey := [fieldparams.BLSPubkeyLength]byte{}
-		copy(nKey[:1], input)
-		return nKey
-	}
-
-	tests := []struct {
-		name            string
-		modifyFunc      func(b *BeaconState, k [fieldparams.BLSPubkeyLength]byte)
-		exists          bool
-		expectedIdx     types.ValidatorIndex
-		largestIdxInSet types.ValidatorIndex
-	}{
-		{
-			name: "retrieve validator",
-			modifyFunc: func(b *BeaconState, key [fieldparams.BLSPubkeyLength]byte) {
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key[:]}))
-			},
-			exists:      true,
-			expectedIdx: 0,
-		},
-		{
-			name: "retrieve validator with multiple validators from the start",
-			modifyFunc: func(b *BeaconState, key [fieldparams.BLSPubkeyLength]byte) {
-				key1 := keyCreator([]byte{'C'})
-				key2 := keyCreator([]byte{'D'})
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key[:]}))
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key1[:]}))
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key2[:]}))
-			},
-			exists:      true,
-			expectedIdx: 0,
-		},
-		{
-			name: "retrieve validator with multiple validators",
-			modifyFunc: func(b *BeaconState, key [fieldparams.BLSPubkeyLength]byte) {
-				key1 := keyCreator([]byte{'C'})
-				key2 := keyCreator([]byte{'D'})
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key1[:]}))
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key2[:]}))
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key[:]}))
-			},
-			exists:      true,
-			expectedIdx: 2,
-		},
-		{
-			name: "retrieve validator with multiple validators from the start with shared state",
-			modifyFunc: func(b *BeaconState, key [fieldparams.BLSPubkeyLength]byte) {
-				key1 := keyCreator([]byte{'C'})
-				key2 := keyCreator([]byte{'D'})
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key[:]}))
-				_ = b.Copy()
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key1[:]}))
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key2[:]}))
-			},
-			exists:      true,
-			expectedIdx: 0,
-		},
-		{
-			name: "retrieve validator with multiple validators with shared state",
-			modifyFunc: func(b *BeaconState, key [fieldparams.BLSPubkeyLength]byte) {
-				key1 := keyCreator([]byte{'C'})
-				key2 := keyCreator([]byte{'D'})
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key1[:]}))
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key2[:]}))
-				n := b.Copy()
-				// Append to another state
-				assert.NoError(t, n.AppendValidator(&ethpb.Validator{PublicKey: key[:]}))
-
-			},
-			exists:      false,
-			expectedIdx: 0,
+func TestBeaconState_MatchCurrentJustifiedCheckpt(t *testing.T) {
+	testtmpl.VerifyBeaconState_MatchCurrentJustifiedCheckptNative(
+		t,
+		func(cp *ethpb.Checkpoint) (state.BeaconState, error) {
+			return InitializeFromProto(&ethpb.BeaconStateAltair{CurrentJustifiedCheckpoint: cp})
 		},
-		{
-			name: "retrieve validator with multiple validators with shared state at boundary",
-			modifyFunc: func(b *BeaconState, key [fieldparams.BLSPubkeyLength]byte) {
-				key1 := keyCreator([]byte{'C'})
-				assert.NoError(t, b.AppendValidator(&ethpb.Validator{PublicKey: key1[:]}))
-				n := b.Copy()
-				// Append to another state
-				assert.NoError(t, n.AppendValidator(&ethpb.Validator{PublicKey: key[:]}))
+	)
+}
 
-			},
-			exists:      false,
-			expectedIdx: 0,
+func TestBeaconState_MatchPreviousJustifiedCheckpt(t *testing.T) {
+	testtmpl.VerifyBeaconState_MatchPreviousJustifiedCheckptNative(
+		t,
+		func(cp *ethpb.Checkpoint) (state.BeaconState, error) {
+			return InitializeFromProto(&ethpb.BeaconStateAltair{PreviousJustifiedCheckpoint: cp})
 		},
-	}
+	)
+}
 
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			s, err := InitializeFromProto(&ethpb.BeaconStateAltair{})
-			require.NoError(t, err)
-			nKey := keyCreator([]byte{'A'})
-			tt.modifyFunc(s, nKey)
-			idx, ok := s.ValidatorIndexByPubkey(nKey)
-			assert.Equal(t, tt.exists, ok)
-			assert.Equal(t, tt.expectedIdx, idx)
-		})
-	}
+func TestBeaconState_ValidatorByPubkey(t *testing.T) {
+	testtmpl.VerifyBeaconState_ValidatorByPubkey(t, func() (state.BeaconState, error) {
+		return InitializeFromProto(&ethpb.BeaconStateAltair{})
+	})
 }

beacon-chain/state/state-native/v2/getters_validator_test.go

@@ -4,18 +4,15 @@ import (
 	"testing"
 
 	"github.com/prysmaticlabs/prysm/beacon-chain/state"
-	v1 "github.com/prysmaticlabs/prysm/beacon-chain/state/state-native/v1"
+	testtmpl "github.com/prysmaticlabs/prysm/beacon-chain/state/testing"
+	v2 "github.com/prysmaticlabs/prysm/beacon-chain/state/v2"
 	ethpb "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1"
-	"github.com/prysmaticlabs/prysm/testing/assert"
-	"github.com/prysmaticlabs/prysm/testing/require"
 )
 
 func TestBeaconState_ValidatorAtIndexReadOnly_HandlesNilSlice(t *testing.T) {
-	st, err := v1.InitializeFromProtoUnsafe(&ethpb.BeaconState{
-		Validators: nil,
+	testtmpl.VerifyBeaconState_ValidatorAtIndexReadOnly_HandlesNilSlice(t, func() (state.BeaconState, error) {
+		return v2.InitializeFromProtoUnsafe(&ethpb.BeaconStateAltair{
+			Validators: nil,
+		})
 	})
-	require.NoError(t, err)
-
-	_, err = st.ValidatorAtIndexReadOnly(0)
-	assert.Equal(t, state.ErrNilValidatorsInState, err)
 }