Merge pull request #6 from nomis52/clock-reset

Add Timer.Reset()

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()
 }
 

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,14 +174,43 @@ 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")
 	case <-time.After(30 * time.Millisecond):
 	}
 }
 
+// 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