feat: add Next() method to mock Clock (#13657)

spikecurtis · web-flow · commit 46e1c36c4200 · 2024-06-25T16:48:26.000+04:00
diff --git a/clock/README.md b/clock/README.md
@@ -188,6 +188,22 @@ w.MustWait(ctx)
 // d contains the duration we advanced
 ```
 
+`d, ok := Peek()` returns the duration until the next event, if any (`ok` is `true`). You can use
+this to advance a specific time, regardless of the tickers and timer events:
+
+```go
+desired := time.Minute // time to advance
+for desired > 0 {
+	p, ok := mClock.Peek()
+	if !ok || p > desired {
+		mClock.Advance(desired).MustWait(ctx)
+		break
+	}
+	mClock.Advance(p).MustWait(ctx)
+	desired -= p
+}
+```
+
 ### Traps
 
 A trap allows you to match specific calls into the library while mocking, block their return,
diff --git a/clock/mock.go b/clock/mock.go
@@ -357,6 +357,18 @@ func (m *Mock) AdvanceNext() (time.Duration, AdvanceWaiter) {
 	return d, w
 }
 
+// Peek returns the duration until the next ticker or timer event and the value
+// true, or, if there are no running tickers or timers, it returns zero and
+// false.
+func (m *Mock) Peek() (d time.Duration, ok bool) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	if m.nextTime.IsZero() {
+		return 0, false
+	}
+	return m.nextTime.Sub(m.cur), true
+}
+
 // Trapper allows the creation of Traps
 type Trapper struct {
 	// mock is the underlying Mock.  This is a thin wrapper around Mock so that
diff --git a/clock/mock_test.go b/clock/mock_test.go
@@ -84,7 +84,7 @@ func TestAfterFunc_NegativeDuration(t *testing.T) {
 func TestNewTicker(t *testing.T) {
 	t.Parallel()
 	// nolint:gocritic // trying to avoid Coder-specific stuff with an eye toward spinning this out
-	ctx, cancel := context.WithTimeout(context.Background(), 1000*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 
 	mClock := clock.NewMock(t)
@@ -167,3 +167,50 @@ func TestNewTicker(t *testing.T) {
 		}
 	}
 }
+
+func TestPeek(t *testing.T) {
+	t.Parallel()
+	// nolint:gocritic // trying to avoid Coder-specific stuff with an eye toward spinning this out
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	mClock := clock.NewMock(t)
+	d, ok := mClock.Peek()
+	if d != 0 {
+		t.Fatal("expected Peek() to return 0")
+	}
+	if ok {
+		t.Fatal("expected Peek() to return false")
+	}
+
+	tmr := mClock.NewTimer(time.Second)
+	d, ok = mClock.Peek()
+	if d != time.Second {
+		t.Fatal("expected Peek() to return 1s")
+	}
+	if !ok {
+		t.Fatal("expected Peek() to return true")
+	}
+
+	mClock.Advance(999 * time.Millisecond).MustWait(ctx)
+	d, ok = mClock.Peek()
+	if d != time.Millisecond {
+		t.Fatal("expected Peek() to return 1ms")
+	}
+	if !ok {
+		t.Fatal("expected Peek() to return true")
+	}
+
+	stopped := tmr.Stop()
+	if !stopped {
+		t.Fatal("expected Stop() to return true")
+	}
+
+	d, ok = mClock.Peek()
+	if d != 0 {
+		t.Fatal("expected Peek() to return 0")
+	}
+	if ok {
+		t.Fatal("expected Peek() to return false")
+	}
+}
