1+ // Package clock provides a mock clock for testing purposes.
12package clock
23
34import (
78 "time"
89)
910
10- // Re -export of time.Duration
11+ // Duration is a re -export of time.Duration.
1112type Duration = time.Duration
1213
1314// Clock represents an interface to the functions in the standard library time
@@ -58,11 +59,13 @@ func (c *clock) Tick(d time.Duration) <-chan time.Time {
5859
5960func (c * clock ) Ticker (d time.Duration ) * Ticker {
6061 t := time .NewTicker (d )
62+
6163 return & Ticker {C : t .C , ticker : t }
6264}
6365
6466func (c * clock ) Timer (d time.Duration ) * Timer {
6567 t := time .NewTimer (d )
68+
6669 return & Timer {C : t .C , timer : t }
6770}
6871
@@ -74,15 +77,12 @@ func (c *clock) WithTimeout(parent context.Context, t time.Duration) (context.Co
7477 return context .WithTimeout (parent , t )
7578}
7679
77- // Mock represents a mock clock that only moves forward programmically .
80+ // Mock represents a mock clock that only moves forward programmatically .
7881// It can be preferable to a real-time clock when testing time-based functionality.
7982type Mock struct {
80- // mu protects all other fields in this struct, and the data that they
81- // point to.
82- mu sync.Mutex
83-
84- now time.Time // current time
85- timers clockTimers // tickers & timers
83+ now time.Time
84+ timers clockTimers
85+ mu sync.Mutex
8686}
8787
8888// NewMock returns an instance of a mock clock.
@@ -100,10 +100,8 @@ func (m *Mock) Add(d time.Duration) {
100100 m .mu .Unlock ()
101101
102102 // Continue to execute timers until there are no more before the new time.
103- for {
104- if ! m .runNextTimer (t ) {
105- break
106- }
103+ for m .runNextTimer (t ) {
104+
107105 }
108106
109107 // Ensure that we end with the new time.
@@ -119,10 +117,8 @@ func (m *Mock) Add(d time.Duration) {
119117// This should only be called from a single goroutine at a time.
120118func (m * Mock ) Set (t time.Time ) {
121119 // Continue to execute timers until there are no more before the new time.
122- for {
123- if ! m .runNextTimer (t ) {
124- break
125- }
120+ for m .runNextTimer (t ) {
121+
126122 }
127123
128124 // Ensure that we end with the new time.
@@ -134,13 +130,14 @@ func (m *Mock) Set(t time.Time) {
134130 gosched ()
135131}
136132
137- // WaitForAllTimers sets the clock until all timers are expired
133+ // WaitForAllTimers sets the clock until all timers are expired.
138134func (m * Mock ) WaitForAllTimers () time.Time {
139135 // Continue to execute timers until there are no more
140136 for {
141137 m .mu .Lock ()
142138 if len (m .timers ) == 0 {
143139 m .mu .Unlock ()
140+
144141 return m .Now ()
145142 }
146143
@@ -154,22 +151,24 @@ func (m *Mock) WaitForAllTimers() time.Time {
154151// runNextTimer executes the next timer in chronological order and moves the
155152// current time to the timer's next tick time. The next time is not executed if
156153// its next time is after the max time. Returns true if a timer was executed.
157- func (m * Mock ) runNextTimer (max time.Time ) bool {
154+ func (m * Mock ) runNextTimer (maxTime time.Time ) bool {
158155 m .mu .Lock ()
159156
160157 // Sort timers by time.
161158 sort .Sort (m .timers )
162159
163- // If we have no more timers then exit.
160+ // If we have no more timers, then exit.
164161 if len (m .timers ) == 0 {
165162 m .mu .Unlock ()
163+
166164 return false
167165 }
168166
169167 // Retrieve next timer. Exit if next tick is after new time.
170168 t := m .timers [0 ]
171- if t .Next ().After (max ) {
169+ if t .Next ().After (maxTime ) {
172170 m .mu .Unlock ()
171+
173172 return false
174173 }
175174
@@ -180,6 +179,7 @@ func (m *Mock) runNextTimer(max time.Time) bool {
180179
181180 // Execute timer.
182181 t .Tick (now )
182+
183183 return true
184184}
185185
@@ -193,6 +193,7 @@ func (m *Mock) After(d time.Duration) <-chan time.Time {
193193func (m * Mock ) AfterFunc (d time.Duration , f func ()) * Timer {
194194 m .mu .Lock ()
195195 defer m .mu .Unlock ()
196+
196197 ch := make (chan time.Time , 1 )
197198 t := & Timer {
198199 c : ch ,
@@ -202,13 +203,15 @@ func (m *Mock) AfterFunc(d time.Duration, f func()) *Timer {
202203 stopped : false ,
203204 }
204205 m .timers = append (m .timers , (* internalTimer )(t ))
206+
205207 return t
206208}
207209
208210// Now returns the current wall time on the mock clock.
209211func (m * Mock ) Now () time.Time {
210212 m .mu .Lock ()
211213 defer m .mu .Unlock ()
214+
212215 return m .now
213216}
214217
@@ -238,6 +241,7 @@ func (m *Mock) Tick(d time.Duration) <-chan time.Time {
238241func (m * Mock ) Ticker (d time.Duration ) * Ticker {
239242 m .mu .Lock ()
240243 defer m .mu .Unlock ()
244+
241245 ch := make (chan time.Time , 1 )
242246 t := & Ticker {
243247 C : ch ,
@@ -247,6 +251,7 @@ func (m *Mock) Ticker(d time.Duration) *Ticker {
247251 next : m .now .Add (d ),
248252 }
249253 m .timers = append (m .timers , (* internalTicker )(t ))
254+
250255 return t
251256}
252257
@@ -265,6 +270,7 @@ func (m *Mock) Timer(d time.Duration) *Timer {
265270 now := m .now
266271 m .mu .Unlock ()
267272 m .runNextTimer (now )
273+
268274 return t
269275}
270276
@@ -276,16 +282,18 @@ func (m *Mock) removeClockTimer(t clockTimer) {
276282 copy (m .timers [i :], m .timers [i + 1 :])
277283 m .timers [len (m .timers )- 1 ] = nil
278284 m .timers = m .timers [:len (m .timers )- 1 ]
285+
279286 break
280287 }
281288 }
289+
282290 sort .Sort (m .timers )
283291}
284292
285293// clockTimer represents an object with an associated start time.
286294type clockTimer interface {
287295 Next () time.Time
288- Tick (time.Time )
296+ Tick (now time.Time )
289297}
290298
291299// clockTimers represents a list of sortable timers.
@@ -318,10 +326,11 @@ func (t *Timer) Stop() bool {
318326 t .mock .removeClockTimer ((* internalTimer )(t ))
319327 t .stopped = true
320328 t .mock .mu .Unlock ()
329+
321330 return registered
322331}
323332
324- // Reset changes the expiry time of the timer
333+ // Reset changes the expiry time of the timer.
325334func (t * Timer ) Reset (d time.Duration ) bool {
326335 if t .timer != nil {
327336 return t .timer .Reset (d )
@@ -333,11 +342,13 @@ func (t *Timer) Reset(d time.Duration) bool {
333342 t .next = t .mock .now .Add (d )
334343
335344 registered := ! t .stopped
345+
336346 if t .stopped {
337347 t .mock .timers = append (t .mock .timers , (* internalTimer )(t ))
338348 }
339349
340350 t .stopped = false
351+
341352 return registered
342353}
343354
@@ -357,6 +368,7 @@ func (t *internalTimer) Tick(now time.Time) {
357368 if t .stopped {
358369 return
359370 }
371+
360372 if t .fn != nil {
361373 // defer function execution until the lock is released, and
362374 defer func () { go t .fn () }()
@@ -368,8 +380,9 @@ func (t *internalTimer) Tick(now time.Time) {
368380 default :
369381 }
370382 }
383+
371384 t .stopped = true
372- t .mock .removeClockTimer (( * internalTimer )( t ) )
385+ t .mock .removeClockTimer (t )
373386}
374387
375388// Ticker holds a channel that receives "ticks" at regular intervals.
@@ -399,6 +412,7 @@ func (t *Ticker) Stop() {
399412func (t * Ticker ) Reset (dur time.Duration ) {
400413 if t .ticker != nil {
401414 t .ticker .Reset (dur )
415+
402416 return
403417 }
404418
@@ -433,13 +447,14 @@ func (t *internalTicker) Tick(now time.Time) {
433447 case t .c <- now :
434448 default :
435449 }
450+
436451 t .next = now .Add (t .d )
437452}
438453
439454// Sleep momentarily so that other goroutines can process.
440455func gosched () { time .Sleep (1 * time .Millisecond ) }
441456
442457var (
443- // type checking
458+ // type checking.
444459 _ Clock = & Mock {}
445460)
0 commit comments