1  package clock
 2  
 3  import (
 4  	"sync"
 5  	"time"
 6  )
 7  
 8  // TestClock can be used in tests to mock time.
 9  type TestClock struct {
10  	currentTime time.Time
11  	timeChanMap map[time.Time][]chan time.Time
12  	timeLock    sync.Mutex
13  	tickSignal  chan time.Duration
14  }
15  
16  // NewTestClock returns a new test clock.
17  func NewTestClock(startTime time.Time) *TestClock {
18  	return &TestClock{
19  		currentTime: startTime,
20  		timeChanMap: make(map[time.Time][]chan time.Time),
21  	}
22  }
23  
24  // NewTestClockWithTickSignal will create a new test clock with an added
25  // channel which will be used to signal when a new ticker is registered.
26  // This is useful when creating a ticker on a separate goroutine and we'd
27  // like to wait for that to happen before advancing the test case.
28  func NewTestClockWithTickSignal(startTime time.Time,
29  	tickSignal chan time.Duration) *TestClock {
30  
31  	testClock := NewTestClock(startTime)
32  	testClock.tickSignal = tickSignal
33  
34  	return testClock
35  }
36  
37  // Now returns the current (test) time.
38  func (c *TestClock) Now() time.Time {
39  	c.timeLock.Lock()
40  	defer c.timeLock.Unlock()
41  
42  	return c.currentTime
43  }
44  
45  // TickAfter returns a channel that will receive a tick after the specified
46  // duration has passed passed by the user set test time.
47  func (c *TestClock) TickAfter(duration time.Duration) <-chan time.Time {
48  	c.timeLock.Lock()
49  	defer func() {
50  		c.timeLock.Unlock()
51  
52  		// Signal that the ticker has been added.
53  		if c.tickSignal != nil {
54  			c.tickSignal <- duration
55  		}
56  	}()
57  
58  	triggerTime := c.currentTime.Add(duration)
59  	ch := make(chan time.Time, 1)
60  
61  	// If already expired, tick immediately.
62  	if !triggerTime.After(c.currentTime) {
63  		ch <- c.currentTime
64  		return ch
65  	}
66  
67  	// Otherwise store the channel until the trigger time is there.
68  	chans := c.timeChanMap[triggerTime]
69  	chans = append(chans, ch)
70  	c.timeChanMap[triggerTime] = chans
71  
72  	return ch
73  }
74  
75  // SetTime sets the (test) time and triggers tick channels when they expire.
76  func (c *TestClock) SetTime(now time.Time) {
77  	c.timeLock.Lock()
78  	defer c.timeLock.Unlock()
79  
80  	c.currentTime = now
81  	remainingChans := make(map[time.Time][]chan time.Time)
82  	for triggerTime, chans := range c.timeChanMap {
83  		// If the trigger time is still in the future, keep this channel
84  		// in the channel map for later.
85  		if triggerTime.After(now) {
86  			remainingChans[triggerTime] = chans
87  			continue
88  		}
89  
90  		for _, c := range chans {
91  			c <- now
92  		}
93  	}
94  
95  	c.timeChanMap = remainingChans
96  }