efritz/glock

View on GitHub
mock.go

Summary

Maintainability
A
0 mins
Test Coverage
A
99%
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
}