mock.go
package glock
import (
"runtime"
"sort"
"sync"
"time"
)
type (
// MockClock is an implementation of Clock that can be moved forward in time
// in increments for testing code that relies on timeouts or other time-sensitive
// constructs.
MockClock struct {
fakeTime time.Time
triggers mockTriggers
tickers mockTickers
afterArgs []time.Duration
tickerArgs []time.Duration
nowLock sync.RWMutex
afterLock sync.RWMutex
tickerLock sync.Mutex
}
mockTrigger struct {
trigger time.Time
ch chan time.Time
}
mockTriggers []*mockTrigger
mockTickers []*mockTicker
)
// Make sure MockClock conforms to the interfaces
var _ Clock = &MockClock{}
// NewMockClock creates a new MockClock with the internal time set
// to time.Now()
func NewMockClock() *MockClock {
return NewMockClockAt(time.Now())
}
// NewMockClockAt creates a new MockClick with the internal time set
// to the provided time.
func NewMockClockAt(now time.Time) *MockClock {
return &MockClock{
fakeTime: now,
tickers: make([]*mockTicker, 0),
afterArgs: make([]time.Duration, 0),
tickerArgs: make([]time.Duration, 0),
}
}
// SetCurrent sets the internal MockClock time to the supplied time.
func (mc *MockClock) SetCurrent(current time.Time) {
mc.nowLock.Lock()
defer mc.nowLock.Unlock()
mc.fakeTime = current
}
// Advance will advance the internal MockClock time by the supplied time.
func (mc *MockClock) Advance(duration time.Duration) {
mc.nowLock.Lock()
now := mc.fakeTime.Add(duration)
mc.fakeTime = now
mc.nowLock.Unlock()
mc.processTriggers(now)
mc.processTickers(now)
}
func (mc *MockClock) processTriggers(now time.Time) {
mc.afterLock.Lock()
defer mc.afterLock.Unlock()
triggered := 0
for _, trigger := range mc.triggers {
if trigger.trigger.Before(now) || trigger.trigger.Equal(now) {
trigger.ch <- trigger.trigger
triggered++
}
}
mc.triggers = mc.triggers[triggered:]
}
func (mc *MockClock) processTickers(now time.Time) {
mc.tickerLock.Lock()
defer mc.tickerLock.Unlock()
for _, ticker := range mc.tickers {
ticker.publish(now)
}
}
// BlockingAdvance will call Advance but only after there is another routine
// which is blocking on the channel result of a call to After.
func (mc *MockClock) BlockingAdvance(duration time.Duration) {
for mc.BlockedOnAfter() == 0 {
runtime.Gosched()
}
mc.Advance(duration)
}
// Now returns the current time internal to the MockClock
func (mc *MockClock) Now() time.Time {
mc.nowLock.RLock()
defer mc.nowLock.RUnlock()
return mc.fakeTime
}
// After returns a channel that will be sent the current internal MockClock
// time once the MockClock's internal time is at or past the provided duration
func (mc *MockClock) After(duration time.Duration) <-chan time.Time {
mc.nowLock.RLock()
triggerTime := mc.fakeTime.Add(duration)
mc.nowLock.RUnlock()
mc.afterLock.Lock()
defer mc.afterLock.Unlock()
trigger := &mockTrigger{
trigger: triggerTime,
ch: make(chan time.Time, 1),
}
mc.triggers = append(mc.triggers, trigger)
mc.afterArgs = append(mc.afterArgs, duration)
sort.Slice(mc.triggers, func(i, j int) bool {
return mc.triggers[i].trigger.Before(mc.triggers[j].trigger)
})
return trigger.ch
}
// BlockedOnAfter returns the number of calls to After that are blocked
// waiting for a call to Advance to trigger them.
func (mc *MockClock) BlockedOnAfter() int {
mc.afterLock.RLock()
defer mc.afterLock.RUnlock()
return len(mc.triggers)
}
// Sleep will block until the internal MockClock time is at or past the
// provided duration
func (mc *MockClock) Sleep(duration time.Duration) {
<-mc.After(duration)
}
// Since returns the time elapsed since t.
func (mc *MockClock) Since(t time.Time) time.Duration {
return mc.Now().Sub(t)
}
// Until returns the duration until t.
func (mc *MockClock) Until(t time.Time) time.Duration {
return t.Sub(mc.Now())
}
// GetAfterArgs returns the duration of each call to After in the
// same order as they were called. The list is cleared each time
// GetAfterArgs is called.
func (mc *MockClock) GetAfterArgs() []time.Duration {
mc.afterLock.Lock()
defer mc.afterLock.Unlock()
args := mc.afterArgs
mc.afterArgs = mc.afterArgs[:0]
return args
}
// GetTickerArgs returns the duration of each call to create a new
// ticker in the same order as they were called. The list is cleared
// each time GetTickerArgs is called.
func (mc *MockClock) GetTickerArgs() []time.Duration {
mc.tickerLock.Lock()
defer mc.tickerLock.Unlock()
args := mc.tickerArgs
mc.tickerArgs = mc.tickerArgs[:0]
return args
}
type mockTicker struct {
clock *MockClock
duration time.Duration
started time.Time
nextTick time.Time
processQueue []time.Time
ch chan time.Time
wakeup chan struct{}
stopped bool
processLock sync.Mutex
stoppedLock sync.RWMutex
}
// NewTicker creates a new Ticker tied to the internal MockClock time that ticks
// at intervals similar to time.NewTicker(). It will also skip or drop ticks
// for slow readers similar to time.NewTicker() as well.
func (mc *MockClock) NewTicker(duration time.Duration) Ticker {
if duration == 0 {
panic("duration cannot be 0")
}
now := mc.Now()
ticker := &mockTicker{
clock: mc,
duration: duration,
started: now,
nextTick: now.Add(duration),
processQueue: make([]time.Time, 0),
ch: make(chan time.Time),
wakeup: make(chan struct{}, 1),
}
mc.tickerLock.Lock()
defer mc.tickerLock.Unlock()
mc.tickers = append(mc.tickers, ticker)
mc.tickerArgs = append(mc.tickerArgs, duration)
go ticker.process()
return ticker
}
// Chan returns a channel which will receive the MockClock's internal time
// at the interval given when creating the ticker.
func (mt *mockTicker) Chan() <-chan time.Time {
return mt.ch
}
// Stop will stop the ticker from ticking
func (mt *mockTicker) Stop() {
mt.stoppedLock.Lock()
defer mt.stoppedLock.Unlock()
mt.stopped = true
mt.wakeup <- struct{}{}
}
func (mt *mockTicker) publish(now time.Time) {
if mt.isStopped() {
return
}
mt.processLock.Lock()
mt.processQueue = append(mt.processQueue, now)
mt.processLock.Unlock()
select {
case mt.wakeup <- struct{}{}:
default:
}
}
func (mt *mockTicker) process() {
defer close(mt.wakeup)
for !mt.isStopped() {
for {
first, ok := mt.pop()
if !ok {
break
}
if mt.nextTick.After(first) {
continue
}
mt.ch <- mt.nextTick
durationMod := first.Sub(mt.started) % mt.duration
if durationMod == 0 {
mt.nextTick = first.Add(mt.duration)
} else if first.Sub(mt.nextTick) > mt.duration {
mt.nextTick = first.Add(mt.duration - durationMod)
} else {
mt.nextTick = mt.nextTick.Add(mt.duration)
}
}
<-mt.wakeup
}
}
func (mt *mockTicker) pop() (time.Time, bool) {
mt.processLock.Lock()
defer mt.processLock.Unlock()
if len(mt.processQueue) == 0 {
return time.Unix(0, 0), false
}
first := mt.processQueue[0]
mt.processQueue = mt.processQueue[1:]
return first, true
}
func (mt *mockTicker) isStopped() bool {
mt.stoppedLock.RLock()
defer mt.stoppedLock.RUnlock()
return mt.stopped
}