add NextWakeup and AdvanceTo functions

clockwork.go

@@ -22,6 +22,19 @@ type FakeClock interface {
 	// Advance advances the FakeClock to a new point in time, ensuring any existing
 	// sleepers are notified appropriately before returning
 	Advance(d time.Duration)
+
+	// AdvanceTo advances the FakeClock to a new point in time, where the time
+	// is specified directly. Returns false if the specified time has allready
+	// passed.
+	AdvanceTo(t time.Time) bool
+
+	// NextWakeup returns the earliest time from now that a blocker would
+	// wait up. Returns a Zero time if there are no more blockers.
+	//
+	// This can be used with Advance to wake subsequent blockers no matter
+	// what time they are blocked on.
+	NextWakeup() time.Time
+
 	// BlockUntil will block until the FakeClock has the given number of
 	// sleepers (callers of Sleep or After)
 	BlockUntil(n int)
@@ -156,12 +169,7 @@ func (fc *fakeClock) NewTicker(d time.Duration) Ticker {
 	return ft
 }
 
-// Advance advances fakeClock to a new point in time, ensuring channels from any
-// previous invocations of After are notified appropriately before returning
-func (fc *fakeClock) Advance(d time.Duration) {
-	fc.l.Lock()
-	defer fc.l.Unlock()
-	end := fc.time.Add(d)
+func (fc *fakeClock) advanceTo(end time.Time) {
 	var newSleepers []*sleeper
 	for _, s := range fc.sleepers {
 		if end.Sub(s.until) >= 0 {
@@ -175,6 +183,25 @@ func (fc *fakeClock) Advance(d time.Duration) {
 	fc.time = end
 }
 
+// Advance advances fakeClock to a new point in time, ensuring channels from any
+// previous invocations of After are notified appropriately before returning
+func (fc *fakeClock) Advance(d time.Duration) {
+	fc.l.Lock()
+	defer fc.l.Unlock()
+	end := fc.time.Add(d)
+	fc.advanceTo(end)
+}
+
+func (fc *fakeClock) AdvanceTo(end time.Time) bool {
+	fc.l.Lock()
+	defer fc.l.Unlock()
+	if fc.time.After(end) {
+		return false
+	}
+	fc.advanceTo(end)
+	return true
+}
+
 // BlockUntil will block until the fakeClock has the given number of sleepers
 // (callers of Sleep or After)
 func (fc *fakeClock) BlockUntil(n int) {
@@ -193,3 +220,17 @@ func (fc *fakeClock) BlockUntil(n int) {
 	fc.l.Unlock()
 	<-b.ch
 }
+
+func (fc *fakeClock) NextWakeup() time.Time {
+	fc.l.Lock()
+	defer fc.l.Unlock()
+
+	var end time.Time
+	for i, s := range fc.sleepers {
+		if i == 0 || s.until.Before(end) {
+			end = s.until
+		}
+	}
+
+	return end
+}

clockwork_test.go

@@ -137,3 +137,32 @@ func TestFakeClockSince(t *testing.T) {
 		t.Fatalf("fakeClock.Since() returned unexpected duration, got: %d, want: %d", fc.Since(now), elapsedTime)
 	}
 }
+
+func TestFakeClockWakeup(t *testing.T) {
+	fc := NewFakeClock()
+
+	if next := fc.NextWakeup(); !next.IsZero() {
+		t.Fatalf("fakeClock.NextWakeup returned an unexpected value: %s", next)
+	}
+
+	_ = fc.After(1 * time.Hour)
+	next := fc.NextWakeup()
+	if got, expected := next.Sub(fc.Now()), 1*time.Hour; got != expected {
+		t.Fatalf("fakeClock.NextWakeup returned %s, expected %s", got, expected)
+	}
+
+	_ = fc.After(20 * time.Minute)
+	next = fc.NextWakeup()
+	if got, expected := next.Sub(fc.Now()), 20*time.Minute; got != expected {
+		t.Fatalf("fakeClock.NextWakeup returned %s, expected %s", got, expected)
+	}
+
+	if ok := fc.AdvanceTo(next); !ok {
+		t.Fatalf("fakeClock.AdvanceTo failed")
+	}
+
+	next = fc.NextWakeup()
+	if got, expected := next.Sub(fc.Now()), 40*time.Minute; got != expected {
+		t.Fatalf("fakeClock.NextWakeup returned %s, expected %s", got, expected)
+	}
+}