fix: allow mock clock Timers to accept negative duration

spikecurtis · spikecurtis · commit d616464318c7 · 2024-06-18T10:48:19.000+04:00
diff --git a/clock/mock.go b/clock/mock.go
@@ -52,9 +52,6 @@ func (m *Mock) TickerFunc(ctx context.Context, d time.Duration, f func() error,
 }
 
 func (m *Mock) NewTimer(d time.Duration, tags ...string) *Timer {
-	if d < 0 {
-		panic("duration must be positive or zero")
-	}
 	m.mu.Lock()
 	defer m.mu.Unlock()
 	c := newCall(clockFunctionNewTimer, tags, withDuration(d))
@@ -67,14 +64,17 @@ func (m *Mock) NewTimer(d time.Duration, tags ...string) *Timer {
 		nxt:  m.cur.Add(d),
 		mock: m,
 	}
+	if d <= 0 {
+		// zero or negative duration timer means we should immediately fire
+		// it, rather than add it.
+		go t.fire(t.mock.cur)
+		return t
+	}
 	m.addTimerLocked(t)
 	return t
 }
 
 func (m *Mock) AfterFunc(d time.Duration, f func(), tags ...string) *Timer {
-	if d < 0 {
-		panic("duration must be positive or zero")
-	}
 	m.mu.Lock()
 	defer m.mu.Unlock()
 	c := newCall(clockFunctionAfterFunc, tags, withDuration(d))
@@ -85,6 +85,12 @@ func (m *Mock) AfterFunc(d time.Duration, f func(), tags ...string) *Timer {
 		fn:   f,
 		mock: m,
 	}
+	if d <= 0 {
+		// zero or negative duration timer means we should immediately fire
+		// it, rather than add it.
+		go t.fire(t.mock.cur)
+		return t
+	}
 	m.addTimerLocked(t)
 	return t
 }
diff --git a/clock/mock_test.go b/clock/mock_test.go
@@ -0,0 +1,82 @@
+package clock_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/coder/coder/v2/clock"
+)
+
+func TestTimer_NegativeDuration(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)
+	start := mClock.Now()
+	trap := mClock.Trap().NewTimer()
+	defer trap.Close()
+
+	timers := make(chan *clock.Timer, 1)
+	go func() {
+		timers <- mClock.NewTimer(-time.Second)
+	}()
+	c := trap.MustWait(ctx)
+	c.Release()
+	// trap returns the actual passed value
+	if c.Duration != -time.Second {
+		t.Fatalf("expected -time.Second, got: %v", c.Duration)
+	}
+
+	tmr := <-timers
+	select {
+	case <-ctx.Done():
+		t.Fatal("timeout waiting for timer")
+	case tme := <-tmr.C:
+		// the tick is the current time, not the past
+		if !tme.Equal(start) {
+			t.Fatalf("expected time %v, got %v", start, tme)
+		}
+	}
+	if tmr.Stop() {
+		t.Fatal("timer still running")
+	}
+}
+
+func TestAfterFunc_NegativeDuration(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)
+	trap := mClock.Trap().AfterFunc()
+	defer trap.Close()
+
+	timers := make(chan *clock.Timer, 1)
+	done := make(chan struct{})
+	go func() {
+		timers <- mClock.AfterFunc(-time.Second, func() {
+			close(done)
+		})
+	}()
+	c := trap.MustWait(ctx)
+	c.Release()
+	// trap returns the actual passed value
+	if c.Duration != -time.Second {
+		t.Fatalf("expected -time.Second, got: %v", c.Duration)
+	}
+
+	tmr := <-timers
+	select {
+	case <-ctx.Done():
+		t.Fatal("timeout waiting for timer")
+	case <-done:
+		// OK!
+	}
+	if tmr.Stop() {
+		t.Fatal("timer still running")
+	}
+}
diff --git a/clock/timer.go b/clock/timer.go
@@ -44,9 +44,6 @@ func (t *Timer) Reset(d time.Duration, tags ...string) bool {
 	if t.timer != nil {
 		return t.timer.Reset(d)
 	}
-	if d < 0 {
-		panic("duration must be positive or zero")
-	}
 	t.mock.mu.Lock()
 	defer t.mock.mu.Unlock()
 	c := newCall(clockFunctionTimerReset, tags, withDuration(d))
@@ -57,9 +54,9 @@ func (t *Timer) Reset(d time.Duration, tags ...string) bool {
 	case <-t.c:
 	default:
 	}
-	if d == 0 {
-		// zero duration timer means we should immediately re-fire it, rather
-		// than remove and re-add it.
+	if d <= 0 {
+		// zero or negative duration timer means we should immediately re-fire
+		// it, rather than remove and re-add it.
 		t.stopped = false
 		go t.fire(t.mock.cur)
 		return result
