/
mock_timer_client.go
95 lines (78 loc) · 1.97 KB
/
mock_timer_client.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
package test
import (
"sync"
"sync/atomic"
"time"
)
type timer struct {
start time.Time
duration time.Duration
repeat bool
fired bool
c chan time.Time
}
// MockTimerClient contains mock timer client.
type MockTimerClient struct {
sync.RWMutex
current time.Time
timers []*timer
afterCalled uint64
tickCalled uint64
}
func (m *MockTimerClient) newTimer(d time.Duration, repeat bool) *timer {
m.RLock()
t := &timer{
start: m.current,
duration: d,
repeat: repeat,
fired: false,
c: make(chan time.Time, 1),
}
m.RUnlock()
m.Lock()
defer m.Unlock()
m.timers = append(m.timers, t)
return t
}
// After is mock of time.After().
func (m *MockTimerClient) After(d time.Duration) <-chan time.Time {
atomic.AddUint64(&m.afterCalled, 1)
return m.newTimer(d, false).c
}
// Tick is mock of time.Tick().
func (m *MockTimerClient) Tick(d time.Duration) <-chan time.Time {
atomic.AddUint64(&m.tickCalled, 1)
return m.newTimer(d, true).c
}
// Advance simulates time passing and signal timers / tickers accordingly
func (m *MockTimerClient) Advance(d time.Duration) {
m.Lock()
m.current = m.current.Add(d)
m.Unlock()
m.RLock()
defer m.RUnlock()
curr := m.current
for _, t := range m.timers {
if t.repeat {
// for Tickers, calculate how many ticks has passed and signal accordingly
for i := int64(0); i < int64(curr.Sub(t.start)/t.duration); i++ {
t.c <- t.start.Add(t.duration * time.Duration(i+1))
t.start = t.start.Add(t.duration)
}
} else {
// for Afters (one-off), signal once
if !t.fired && (curr.Sub(t.start) >= t.duration) {
t.c <- t.start.Add(t.duration)
t.fired = true
}
}
}
}
// AfterCalledTimes calculates number of times after is called.
func (m *MockTimerClient) AfterCalledTimes() uint64 {
return atomic.LoadUint64(&m.afterCalled)
}
// TickCalledTimes calculates number of times tick is called.
func (m *MockTimerClient) TickCalledTimes() uint64 {
return atomic.LoadUint64(&m.tickCalled)
}