From 3c2693f9fbadb9b3e450878665a1a707213fdf56 Mon Sep 17 00:00:00 2001
From: Louis des Landes <louis.deslandes@fleet.space>
Date: Tue, 28 Jan 2025 16:43:20 +1030
Subject: [PATCH] Add NextWakeup to fake clock

So we can find the next time that a clock is due to wake for
easy advancing of the clock.
---
 clockwork.go      | 18 ++++++++++++++++++
 clockwork_test.go | 27 +++++++++++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/clockwork.go b/clockwork.go
index 85a9934..e16a7c8 100644
--- a/clockwork.go
+++ b/clockwork.go
@@ -317,3 +317,21 @@ func (fc *FakeClock) setExpirer(e expirer, d time.Duration) {
 		return false
 	})
 }
+
+// NextWakeup returns the time of the next expirer to wake up.
+//
+// If there are no waiters, it returns the zero time.
+//
+// This can be used with Advance to wake subsequent blockers no matter
+// what time they are blocked on.
+func (fc *FakeClock) NextWakeup() time.Time {
+	fc.l.RLock()
+	defer fc.l.RUnlock()
+	var end time.Time
+	for i, s := range fc.waiters {
+		if i == 0 || s.expiration().Before(end) {
+			end = s.expiration()
+		}
+	}
+	return end
+}
diff --git a/clockwork_test.go b/clockwork_test.go
index 7f25606..5c18543 100644
--- a/clockwork_test.go
+++ b/clockwork_test.go
@@ -209,3 +209,30 @@ func TestFakeClockRace(t *testing.T) {
 	go func() { fc.NewTimer(d) }()
 	go func() { fc.Sleep(d) }()
 }
+
+func TestNextWakeup(t *testing.T) {
+	t.Parallel()
+	zero := time.Time{}
+	fc := NewFakeClockAt(zero)
+	if got := fc.NextWakeup(); got != zero {
+		t.Errorf("fakeClock.nextWakeup() returned %v, want zero", got)
+	}
+
+	ch := fc.After(time.Minute)
+
+	if got := fc.NextWakeup(); got != zero.Add(time.Minute) {
+		t.Errorf("fakeClock.nextWakeup() returned %v, want 1 minute", got)
+	}
+
+	fc.Advance(fc.NextWakeup().Sub(fc.Now()))
+	<-ch
+
+	if got := fc.NextWakeup(); got != zero {
+		t.Errorf("fakeClock.nextWakeup() returned %v, want zero", got)
+	}
+
+	// Check now
+	if got := fc.Now(); got != zero.Add(time.Minute) {
+		t.Errorf("fakeClock.Now() returned %v, want zero+1 minute", got)
+	}
+}