fix scheduler tests

ssvlabs · Dec 3, 2024 · 26de334 · 26de334
1 parent 160c6c5
commit 26de334
Show file tree

Hide file tree

Showing 5 changed files with 174 additions and 0 deletions.
diff --git a/operator/duties/attester_genesis_test.go b/operator/duties/attester_genesis_test.go
@@ -1,6 +1,7 @@
 package duties
 
 import (
+	"bytes"
 	"context"
 	"testing"
 	"time"
@@ -58,6 +59,35 @@ func setupAttesterGenesisDutiesMock(
 	}
 	s.validatorProvider.(*MockValidatorProvider).EXPECT().SelfParticipatingValidators(gomock.Any()).DoAndReturn(getShares).AnyTimes()
 	s.validatorProvider.(*MockValidatorProvider).EXPECT().ParticipatingValidators(gomock.Any()).DoAndReturn(getShares).AnyTimes()
+	s.validatorProvider.(*MockValidatorProvider).EXPECT().Validator(gomock.Any()).DoAndReturn(
+		func(pubKey []byte) (*types.SSVShare, bool) {
+			var ssvShare *types.SSVShare
+			var minEpoch phase0.Epoch
+			dutiesMap.Range(func(epoch phase0.Epoch, duties []*eth2apiv1.AttesterDuty) bool {
+				for _, duty := range duties {
+					if bytes.Equal(duty.PubKey[:], pubKey) {
+						ssvShare = &types.SSVShare{
+							Share: spectypes.Share{
+								ValidatorIndex: duty.ValidatorIndex,
+							},
+						}
+						if epoch < minEpoch {
+							minEpoch = epoch
+							ssvShare.SetMinParticipationEpoch(epoch)
+						}
+						return true
+					}
+				}
+				return true
+			})
+
+			if ssvShare != nil {
+				return ssvShare, true
+			}
+
+			return nil, false
+		},
+	).AnyTimes()
 
 	s.beaconNode.(*MockBeaconNode).EXPECT().SubmitBeaconCommitteeSubscriptions(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 

diff --git a/operator/duties/attester_test.go b/operator/duties/attester_test.go
@@ -1,6 +1,7 @@
 package duties
 
 import (
+	"bytes"
 	"context"
 	"testing"
 	"time"
@@ -56,6 +57,35 @@ func setupAttesterDutiesMock(
 	}
 	s.validatorProvider.(*MockValidatorProvider).EXPECT().SelfParticipatingValidators(gomock.Any()).DoAndReturn(getShares).AnyTimes()
 	s.validatorProvider.(*MockValidatorProvider).EXPECT().ParticipatingValidators(gomock.Any()).DoAndReturn(getShares).AnyTimes()
+	s.validatorProvider.(*MockValidatorProvider).EXPECT().Validator(gomock.Any()).DoAndReturn(
+		func(pubKey []byte) (*types.SSVShare, bool) {
+			var ssvShare *types.SSVShare
+			var minEpoch phase0.Epoch
+			dutiesMap.Range(func(epoch phase0.Epoch, duties []*eth2apiv1.AttesterDuty) bool {
+				for _, duty := range duties {
+					if bytes.Equal(duty.PubKey[:], pubKey) {
+						ssvShare = &types.SSVShare{
+							Share: spectypes.Share{
+								ValidatorIndex: duty.ValidatorIndex,
+							},
+						}
+						if epoch < minEpoch {
+							minEpoch = epoch
+							ssvShare.SetMinParticipationEpoch(epoch)
+						}
+						return true
+					}
+				}
+				return true
+			})
+
+			if ssvShare != nil {
+				return ssvShare, true
+			}
+
+			return nil, false
+		},
+	).AnyTimes()
 
 	s.beaconNode.(*MockBeaconNode).EXPECT().SubmitBeaconCommitteeSubscriptions(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 

diff --git a/operator/duties/committee_test.go b/operator/duties/committee_test.go
@@ -1,6 +1,7 @@
 package duties
 
 import (
+	"bytes"
 	"context"
 	"testing"
 	"time"
@@ -76,6 +77,57 @@ func setupCommitteeDutiesMock(
 
 	s.validatorProvider.(*MockValidatorProvider).EXPECT().SelfParticipatingValidators(gomock.Any()).Return(activeShares).AnyTimes()
 	s.validatorProvider.(*MockValidatorProvider).EXPECT().ParticipatingValidators(gomock.Any()).Return(activeShares).AnyTimes()
+	s.validatorProvider.(*MockValidatorProvider).EXPECT().Validator(gomock.Any()).DoAndReturn(
+		func(pubKey []byte) (*ssvtypes.SSVShare, bool) {
+			var ssvShare *ssvtypes.SSVShare
+			var minEpoch phase0.Epoch
+			attDuties.Range(func(epoch phase0.Epoch, duties []*eth2apiv1.AttesterDuty) bool {
+				for _, duty := range duties {
+					if bytes.Equal(duty.PubKey[:], pubKey) {
+						ssvShare = &ssvtypes.SSVShare{
+							Share: spectypes.Share{
+								ValidatorIndex: duty.ValidatorIndex,
+							},
+						}
+						if epoch < minEpoch {
+							minEpoch = epoch
+							ssvShare.SetMinParticipationEpoch(epoch)
+						}
+						return true
+					}
+				}
+				return true
+			})
+
+			if ssvShare == nil {
+				minEpoch = phase0.Epoch(0)
+				syncDuties.Range(func(period uint64, duties []*eth2apiv1.SyncCommitteeDuty) bool {
+					for _, duty := range duties {
+						if bytes.Equal(duty.PubKey[:], pubKey) {
+							ssvShare = &ssvtypes.SSVShare{
+								Share: spectypes.Share{
+									ValidatorIndex: duty.ValidatorIndex,
+								},
+							}
+							firstEpoch := s.network.Beacon.FirstEpochOfSyncPeriod(period)
+							if firstEpoch < minEpoch {
+								minEpoch = firstEpoch
+								ssvShare.SetMinParticipationEpoch(firstEpoch)
+							}
+							return true
+						}
+					}
+					return true
+				})
+			}
+
+			if ssvShare != nil {
+				return ssvShare, true
+			}
+
+			return nil, false
+		},
+	).AnyTimes()
 
 	s.validatorController.(*MockValidatorController).EXPECT().AllActiveIndices(gomock.Any(), gomock.Any()).DoAndReturn(
 		func(epoch phase0.Epoch, afterInit bool) []phase0.ValidatorIndex {

diff --git a/operator/duties/sync_committee_genesis_test.go b/operator/duties/sync_committee_genesis_test.go
@@ -1,6 +1,7 @@
 package duties
 
 import (
+	"bytes"
 	"context"
 	"testing"
 	"time"
@@ -69,6 +70,36 @@ func setupSyncCommitteeGenesisDutiesMock(
 
 	s.validatorProvider.(*MockValidatorProvider).EXPECT().SelfParticipatingValidators(gomock.Any()).Return(activeShares).AnyTimes()
 	s.validatorProvider.(*MockValidatorProvider).EXPECT().ParticipatingValidators(gomock.Any()).Return(activeShares).AnyTimes()
+	s.validatorProvider.(*MockValidatorProvider).EXPECT().Validator(gomock.Any()).DoAndReturn(
+		func(pubKey []byte) (*ssvtypes.SSVShare, bool) {
+			var ssvShare *ssvtypes.SSVShare
+			var minEpoch phase0.Epoch
+			dutiesMap.Range(func(period uint64, duties []*v1.SyncCommitteeDuty) bool {
+				for _, duty := range duties {
+					if bytes.Equal(duty.PubKey[:], pubKey) {
+						ssvShare = &ssvtypes.SSVShare{
+							Share: spectypes.Share{
+								ValidatorIndex: duty.ValidatorIndex,
+							},
+						}
+						firstEpoch := s.network.Beacon.FirstEpochOfSyncPeriod(period)
+						if firstEpoch < minEpoch {
+							minEpoch = firstEpoch
+							ssvShare.SetMinParticipationEpoch(firstEpoch)
+						}
+						return true
+					}
+				}
+				return true
+			})
+
+			if ssvShare != nil {
+				return ssvShare, true
+			}
+
+			return nil, false
+		},
+	).AnyTimes()
 
 	s.validatorController.(*MockValidatorController).EXPECT().AllActiveIndices(gomock.Any(), gomock.Any()).DoAndReturn(
 		func(epoch phase0.Epoch, afterInit bool) []phase0.ValidatorIndex {

diff --git a/operator/duties/sync_committee_test.go b/operator/duties/sync_committee_test.go
@@ -1,6 +1,7 @@
 package duties
 
 import (
+	"bytes"
 	"context"
 	"testing"
 	"time"
@@ -67,6 +68,36 @@ func setupSyncCommitteeDutiesMock(
 
 	s.validatorProvider.(*MockValidatorProvider).EXPECT().SelfParticipatingValidators(gomock.Any()).Return(activeShares).AnyTimes()
 	s.validatorProvider.(*MockValidatorProvider).EXPECT().ParticipatingValidators(gomock.Any()).Return(activeShares).AnyTimes()
+	s.validatorProvider.(*MockValidatorProvider).EXPECT().Validator(gomock.Any()).DoAndReturn(
+		func(pubKey []byte) (*ssvtypes.SSVShare, bool) {
+			var ssvShare *ssvtypes.SSVShare
+			var minEpoch phase0.Epoch
+			dutiesMap.Range(func(period uint64, duties []*v1.SyncCommitteeDuty) bool {
+				for _, duty := range duties {
+					if bytes.Equal(duty.PubKey[:], pubKey) {
+						ssvShare = &ssvtypes.SSVShare{
+							Share: spectypes.Share{
+								ValidatorIndex: duty.ValidatorIndex,
+							},
+						}
+						firstEpoch := s.network.Beacon.FirstEpochOfSyncPeriod(period)
+						if firstEpoch < minEpoch {
+							minEpoch = firstEpoch
+							ssvShare.SetMinParticipationEpoch(firstEpoch)
+						}
+						return true
+					}
+				}
+				return true
+			})
+
+			if ssvShare != nil {
+				return ssvShare, true
+			}
+
+			return nil, false
+		},
+	).AnyTimes()
 
 	s.validatorController.(*MockValidatorController).EXPECT().AllActiveIndices(gomock.Any(), gomock.Any()).DoAndReturn(
 		func(epoch phase0.Epoch, afterInit bool) []phase0.ValidatorIndex {