diff --git a/clock.go b/clock.go index 3d1d3d5..518178d 100644 --- a/clock.go +++ b/clock.go @@ -192,10 +192,11 @@ func (m *Mock) Timer(d time.Duration) *Timer { defer m.mu.Unlock() ch := make(chan time.Time, 1) t := &Timer{ - C: ch, - c: ch, - mock: m, - next: m.now.Add(d), + C: ch, + c: ch, + mock: m, + next: m.now.Add(d), + stopped: false, } m.timers = append(m.timers, (*internalTimer)(t)) return t @@ -231,21 +232,42 @@ func (a clockTimers) Less(i, j int) bool { return a[i].Next().Before(a[j].Next() // Timer represents a single event. // The current time will be sent on C, unless the timer was created by AfterFunc. type Timer struct { - C <-chan time.Time - c chan time.Time - timer *time.Timer // realtime impl, if set - next time.Time // next tick time - mock *Mock // mock clock, if set - fn func() // AfterFunc function, if set + C <-chan time.Time + c chan time.Time + timer *time.Timer // realtime impl, if set + next time.Time // next tick time + mock *Mock // mock clock, if set + fn func() // AfterFunc function, if set + stopped bool // True if stopped, false if running } // Stop turns off the ticker. -func (t *Timer) Stop() { +func (t *Timer) Stop() bool { if t.timer != nil { - t.timer.Stop() - } else { - t.mock.removeClockTimer((*internalTimer)(t)) + return t.timer.Stop() + } + + registered := !t.stopped + t.mock.removeClockTimer((*internalTimer)(t)) + t.stopped = true + return registered +} + +// Reset changes the expiry time of the timer +func (t *Timer) Reset(d time.Duration) bool { + if t.timer != nil { + return t.timer.Reset(d) + } + + t.next = t.mock.now.Add(d) + registered := !t.stopped + if t.stopped { + t.mock.mu.Lock() + t.mock.timers = append(t.mock.timers, (*internalTimer)(t)) + t.mock.mu.Unlock() } + t.stopped = false + return registered } type internalTimer Timer @@ -258,6 +280,7 @@ func (t *internalTimer) Tick(now time.Time) { t.c <- now } t.mock.removeClockTimer((*internalTimer)(t)) + t.stopped = true gosched() } diff --git a/clock_test.go b/clock_test.go index 4101ab9..9c090d5 100644 --- a/clock_test.go +++ b/clock_test.go @@ -159,6 +159,10 @@ func TestClock_Timer(t *testing.T) { if !ok { t.Fatal("too early") } + + if timer.Stop() { + t.Fatal("timer still running") + } } // Ensure that the clock's timer can be stopped. @@ -170,7 +174,12 @@ func TestClock_Timer_Stop(t *testing.T) { }() timer := New().Timer(20 * time.Millisecond) - timer.Stop() + if !timer.Stop() { + t.Fatal("timer not running") + } + if timer.Stop() { + t.Fatal("timer wasn't cancelled") + } select { case <-timer.C: t.Fatal("unexpected send") @@ -178,6 +187,30 @@ func TestClock_Timer_Stop(t *testing.T) { } } +// Ensure that the clock's timer can be reset. +func TestClock_Timer_Reset(t *testing.T) { + var ok bool + go func() { + time.Sleep(20 * time.Millisecond) + ok = true + }() + go func() { + time.Sleep(30 * time.Millisecond) + t.Fatal("too late") + }() + gosched() + + timer := New().Timer(10 * time.Millisecond) + if !timer.Reset(20 * time.Millisecond) { + t.Fatal("timer not running") + } + + <-timer.C + if !ok { + t.Fatal("too early") + } +} + // Ensure that the mock's After channel sends at the correct time. func TestMock_After(t *testing.T) { var ok int32