diff --git a/validator/client/mock_validator.go b/validator/client/mock_validator.go index afe7c5785..6ccd20e3d 100644 --- a/validator/client/mock_validator.go +++ b/validator/client/mock_validator.go @@ -43,6 +43,10 @@ type FakeValidator struct { PubkeysToStatusesMap map[[48]byte]ethpb.ValidatorStatus } +type ctxKey string + +var allValidatorsAreExitedCtxKey = ctxKey("exited") + // Done for mocking. func (fv *FakeValidator) Done() { fv.DoneCalled = true @@ -178,6 +182,9 @@ func (fv *FakeValidator) PubkeysToStatuses(_ context.Context) map[[48]byte]ethpb } // AllValidatorsAreExited for mocking -func (fv *FakeValidator) AllValidatorsAreExited(_ context.Context) (bool, error) { - return false, nil +func (fv *FakeValidator) AllValidatorsAreExited(ctx context.Context) (bool, error) { + if ctx.Value(allValidatorsAreExitedCtxKey) == nil { + return false, nil + } + return ctx.Value(allValidatorsAreExitedCtxKey).(bool), nil } diff --git a/validator/client/runner.go b/validator/client/runner.go index cbe994815..66704259a 100644 --- a/validator/client/runner.go +++ b/validator/client/runner.go @@ -86,9 +86,6 @@ func run(ctx context.Context, v Validator) { handleAssignmentError(err, headSlot) } - // We initially assume not all validators are exited - allExited := false - for { ctx, span := trace.StartSpan(ctx, "validator.processSlot") @@ -100,15 +97,12 @@ func run(ctx context.Context, v Validator) { case slot := <-v.NextSlot(): span.AddAttributes(trace.Int64Attribute("slot", int64(slot))) - if allExited { - log.Info("All validators are exited, no more work to perform...") - continue - } - allExited, err = v.AllValidatorsAreExited(ctx) + allExited, err := v.AllValidatorsAreExited(ctx) if err != nil { log.WithError(err).Error("Could not check if validators are exited") } if allExited { + log.Info("All validators are exited, no more work to perform...") continue } diff --git a/validator/client/runner_test.go b/validator/client/runner_test.go index 90481f0b0..fd253839c 100644 --- a/validator/client/runner_test.go +++ b/validator/client/runner_test.go @@ -177,3 +177,20 @@ func TestBothProposesAndAttests_NextSlot(t *testing.T) { require.Equal(t, true, v.ProposeBlockCalled, "ProposeBlock(%d) was not called", slot) assert.Equal(t, slot, v.ProposeBlockArg1, "ProposeBlock was called with wrong arg") } + +func TestAllValidatorsAreExited_NextSlot(t *testing.T) { + v := &FakeValidator{} + ctx, cancel := context.WithCancel(context.WithValue(context.Background(), allValidatorsAreExitedCtxKey, true)) + hook := logTest.NewGlobal() + + slot := uint64(55) + ticker := make(chan uint64) + v.NextSlotRet = ticker + go func() { + ticker <- slot + + cancel() + }() + run(ctx, v) + assert.LogsContain(t, hook, "All validators are exited") +}