From 655fb0dec44a9e5ae9f2bf9f26a4a3fd09413ff3 Mon Sep 17 00:00:00 2001
From: xenowits <kumar.abhishek111b@gmail.com>
Date: Wed, 20 Jul 2022 14:45:23 +0530
Subject: [PATCH] fix tests

---
 core/deadline.go      | 18 +++++------
 core/deadline_test.go | 72 ++++++++++++++++++++++++++++++-------------
 2 files changed, 59 insertions(+), 31 deletions(-)

diff --git a/core/deadline.go b/core/deadline.go
index 87329bd97..c80a774a2 100644
--- a/core/deadline.go
+++ b/core/deadline.go
@@ -38,8 +38,9 @@ type slotTimeProvider interface {
 // may only be used by a single goroutine. So, multiple instances are required
 // for different components and use cases.
 type Deadliner interface {
-	// Add adds a duty to be notified of the Deadline via C. It returns true if the duty is added successfully.
-	// Note that duties will be deduplicated and only a single duty will be provided via C.
+	// Add returns true if the duty was added for future deadline scheduling. It is idempotent
+	// and returns true if the duty was previously added and still awaits deadline scheduling. It
+	// returns false if the duty has already expired and cannot therefore be added for scheduling.
 	Add(duty Duty) bool
 
 	// C returns the same read channel every time and contains deadlined duties.
@@ -84,7 +85,7 @@ func NewDeadliner(ctx context.Context, deadlineFunc func(Duty) time.Time) *Deadl
 		}
 
 		// TODO(dhruv): optimise getCurrDuty and updating current state if earlier deadline detected,
-		// using min heap or ordered map
+		//  using min heap or ordered map
 		for {
 			select {
 			case <-ctx.Done():
@@ -95,7 +96,7 @@ func NewDeadliner(ctx context.Context, deadlineFunc func(Duty) time.Time) *Deadl
 				// Ignore expired duties.
 				if deadline.Before(time.Now()) {
 					d.addedChan <- false
-					return
+					continue
 				}
 
 				duties[duty] = true
@@ -116,18 +117,15 @@ func NewDeadliner(ctx context.Context, deadlineFunc func(Duty) time.Time) *Deadl
 	return d
 }
 
-// Add adds a duty to be notified of the deadline.
+// Add adds a duty to be notified of the deadline. It returns true if the duty
+// was added successfully.
 func (d *Deadline) Add(duty Duty) bool {
 	select {
 	case <-d.quit:
 		return false
 	default:
 		d.dutyChan <- duty
-	}
-
-	select {
-	case ok := <-d.addedChan:
-		return ok
+		return <-d.addedChan
 	}
 }
 
diff --git a/core/deadline_test.go b/core/deadline_test.go
index b5436e408..3fc564e84 100644
--- a/core/deadline_test.go
+++ b/core/deadline_test.go
@@ -29,43 +29,73 @@ func TestDeadliner(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
+	expiredDuties, nonExpiredDuties, voluntaryExits, dutyExpired := setupData(t)
+
 	deadlineFuncProvider := func() func(duty core.Duty) time.Time {
+		startTime := time.Now()
 		return func(duty core.Duty) time.Time {
-			if duty.Type == core.DutyExit {
-				// Do not timeout exit duties.
-				return time.Date(9999, 1, 1, 0, 0, 0, 0, time.UTC)
+			if dutyExpired(duty) || (duty.Type == core.DutyExit) {
+				return startTime.Add(-1 * time.Hour)
 			}
 
-			return time.Now().Add(time.Millisecond * time.Duration(duty.Slot))
+			return startTime.Add(time.Second)
 		}
 	}
 
 	deadliner := core.NewDeadliner(ctx, deadlineFuncProvider())
 
-	expectedDuties := []core.Duty{
-		core.NewVoluntaryExit(2),
-		core.NewAttesterDuty(2),
+	// Add our duties to the deadliner.
+	addDuties(t, expiredDuties, false, deadliner)
+	addDuties(t, nonExpiredDuties, true, deadliner)
+	addDuties(t, voluntaryExits, false, deadliner)
+
+	for i := 0; i < len(nonExpiredDuties); i++ {
+		actualDuty := <-deadliner.C()
+		require.Equal(t, actualDuty, nonExpiredDuties[i])
+	}
+}
+
+// sendDuties runs a goroutine which adds the duties to the deadliner channel.
+func addDuties(t *testing.T, duties []core.Duty, expected bool, deadliner *core.Deadline) {
+	t.Helper()
+
+	go func(duties []core.Duty, expected bool) {
+		for _, duty := range duties {
+			require.Equal(t, deadliner.Add(duty), expected)
+		}
+	}(duties, expected)
+}
+
+// setupData sets up the duties to send to deadliner.
+func setupData(t *testing.T) ([]core.Duty, []core.Duty, []core.Duty, func(core.Duty) bool) {
+	t.Helper()
+
+	expiredDuties := []core.Duty{
 		core.NewAttesterDuty(1),
-		core.NewAttesterDuty(3),
+		core.NewProposerDuty(2),
+		core.NewRandaoDuty(3),
 	}
 
-	for _, duty := range expectedDuties {
-		require.Equal(t, deadliner.Add(duty), true)
+	nonExpiredDuties := []core.Duty{
+		core.NewProposerDuty(1),
+		core.NewAttesterDuty(2),
+		core.NewBuilderProposerDuty(3),
 	}
 
-	var actualDuties []core.Duty
-	for i := 0; i < len(expectedDuties)-1; i++ {
-		actualDuty := <-deadliner.C()
-		actualDuties = append(actualDuties, actualDuty)
+	voluntaryExits := []core.Duty{
+		core.NewVoluntaryExit(2),
+		core.NewVoluntaryExit(4),
 	}
 
-	require.Equal(t, len(expectedDuties), len(actualDuties)+1)
+	dutyExpired := func(duty core.Duty) bool {
+		for _, d := range expiredDuties {
+			if d == duty {
+				return true
+			}
+		}
 
-	// Since DutyExit doesn't timeout, we won't receive it from the deadliner.
-	require.NotEqual(t, expectedDuties[0], actualDuties[0])
+		return false
+	}
 
-	// AttesterDuty for Slot 1 times out before AttesterDuty for Slot 2
-	require.Equal(t, expectedDuties[2], actualDuties[0])
-	require.Equal(t, expectedDuties[1], actualDuties[1])
-	require.Equal(t, expectedDuties[3], actualDuties[2])
+	return expiredDuties, nonExpiredDuties, voluntaryExits, dutyExpired
 }