This repository has been archived by the owner on Oct 11, 2022. It is now read-only.
/
mock.go
132 lines (107 loc) · 3.45 KB
/
mock.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
package clock
import (
"fmt"
"sync"
"time"
)
// MockClock provides a Clock whose time only changes or advances when
// manually specified to do so. This is useful for unit tests.
type MockClock struct {
// Mutex is used to guard the current time, and the cond is used to
// broadcast changes to things waiting for the time to change.
cond *sync.Cond
now time.Time
}
var _ Clock = new(MockClock)
// Now returns the current local time.
func (m *MockClock) Now() time.Time {
m.cond.L.Lock()
defer m.cond.L.Unlock()
return m.now
}
// After waits for the duration to elapse and then sends the current time on the returned channel.
func (m *MockClock) After(d time.Duration) <-chan time.Time {
ch := make(chan time.Time, 1)
target := m.Now().Add(d)
go func() {
for {
m.cond.L.Lock()
if !target.After(m.now) {
now := m.now
m.cond.L.Unlock()
ch <- now
return
}
m.cond.Wait()
m.cond.L.Unlock()
}
}()
return ch
}
// Sleep pauses the current goroutine for at least the duration d. A negative or zero duration causes Sleep to return immediately.
func (m *MockClock) Sleep(d time.Duration) {
<-m.After(d)
}
// Tick is a convenience wrapper for NewTicker providing access to the ticking channel only. While Tick is useful for clients that have no need to shut down the Ticker, be aware that without a way to shut it down the underlying Ticker cannot be recovered by the garbage collector; it "leaks".
func (m *MockClock) Tick(d time.Duration) <-chan time.Time {
return m.NewTicker(d).Chan()
}
// AfterFunc waits for the duration to elapse and then calls f in its own goroutine. It returns a Timer that can be used to cancel the call using its Stop method.
func (m *MockClock) AfterFunc(d time.Duration, f func()) Timer {
t := m.NewTimer(d)
go func() {
<-t.Chan()
go f()
}()
return t
}
// NewTimer creates a new mock Timer that will send the current time on its channel after at least duration d.
func (m *MockClock) NewTimer(d time.Duration) Timer {
t := NewMockTimer(m)
t.Reset(d)
return t
}
// NewTicker returns a new mock Ticker containing a channel that will send the time with a period specified by the duration argument.
// Note: unlike the default ticker included in Go, the mock ticker will *never* skip ticks as time advances.
func (m *MockClock) NewTicker(d time.Duration) Ticker {
return NewMockTicker(m, d)
}
// Since returns the time elapsed since t.
func (m *MockClock) Since(t time.Time) time.Duration {
return m.Now().Sub(t)
}
// SetTime sets the mock clock's time to the given absolute time.
func (m *MockClock) SetTime(t time.Time) {
m.cond.L.Lock()
defer m.cond.L.Unlock()
assertFuture(m.now, t)
m.now = t
m.cond.Broadcast()
}
// AddTime adds the given time duration to the clock.
func (m *MockClock) AddTime(d time.Duration) {
m.cond.L.Lock()
defer m.cond.L.Unlock()
assertFuture(m.now, m.now.Add(d))
m.now = m.now.Add(d)
m.cond.Broadcast()
}
func assertFuture(a, b time.Time) {
if a.After(b) {
panic(fmt.Sprintf("Tried to tick backwards from %v to %v, but cannot travel into the past!", a, b))
}
}
// NewMockClock creates a new mock clock, with its current time set to the provided
// optional start time.
func NewMockClock(start ...time.Time) *MockClock {
m := &MockClock{cond: sync.NewCond(new(sync.Mutex))}
if len(start) > 1 {
panic(fmt.Sprintf("Expected one argument to clock.NewMock, got %d", len(start)))
}
if len(start) == 1 {
m.SetTime(start[0])
return m
}
m.SetTime(time.Now().UTC())
return m
}