github.com/gofunct/common@v0.0.0-20190131174352-fd058c7fbf22/pkg/clock/testing/fake_clock.go

package testing

import (
	"sync"
	"time"

	"k8s.io/utils/clock"
)

var (
	_ = clock.Clock(&FakeClock{})
	_ = clock.Clock(&IntervalClock{})
)

// FakeClock implements clock.Clock, but returns an arbitrary time.
type FakeClock struct {
	lock sync.RWMutex
	time time.Time

	// waiters are waiting for the fake time to pass their specified time
	waiters []*fakeClockWaiter
}

type fakeClockWaiter struct {
	targetTime    time.Time
	stepInterval  time.Duration
	skipIfBlocked bool
	destChan      chan time.Time
	fired         bool
}

// NewFakeClock constructs a fake clock set to the provided time.
func NewFakeClock(t time.Time) *FakeClock {
	return &FakeClock{
		time: t,
	}
}

// Now returns f's time.
func (f *FakeClock) Now() time.Time {
	f.lock.RLock()
	defer f.lock.RUnlock()
	return f.time
}

// Since returns time since the time in f.
func (f *FakeClock) Since(ts time.Time) time.Duration {
	f.lock.RLock()
	defer f.lock.RUnlock()
	return f.time.Sub(ts)
}

// After is the fake version of time.After(d).
func (f *FakeClock) After(d time.Duration) <-chan time.Time {
	f.lock.Lock()
	defer f.lock.Unlock()
	stopTime := f.time.Add(d)
	ch := make(chan time.Time, 1) // Don't block!
	f.waiters = append(f.waiters, &fakeClockWaiter{
		targetTime: stopTime,
		destChan:   ch,
	})
	return ch
}

// NewTimer constructs a fake timer, akin to time.NewTimer(d).
func (f *FakeClock) NewTimer(d time.Duration) clock.Timer {
	f.lock.Lock()
	defer f.lock.Unlock()
	stopTime := f.time.Add(d)
	ch := make(chan time.Time, 1) // Don't block!
	timer := &fakeTimer{
		fakeClock: f,
		waiter: fakeClockWaiter{
			targetTime: stopTime,
			destChan:   ch,
		},
	}
	f.waiters = append(f.waiters, &timer.waiter)
	return timer
}

// Tick constructs a fake ticker, akin to time.Tick
func (f *FakeClock) Tick(d time.Duration) <-chan time.Time {
	if d <= 0 {
		return nil
	}
	f.lock.Lock()
	defer f.lock.Unlock()
	tickTime := f.time.Add(d)
	ch := make(chan time.Time, 1) // hold one tick
	f.waiters = append(f.waiters, &fakeClockWaiter{
		targetTime:    tickTime,
		stepInterval:  d,
		skipIfBlocked: true,
		destChan:      ch,
	})

	return ch
}

// Step moves the clock by Duration and notifies anyone that's called After,
// Tick, or NewTimer.
func (f *FakeClock) Step(d time.Duration) {
	f.lock.Lock()
	defer f.lock.Unlock()
	f.setTimeLocked(f.time.Add(d))
}

// SetTime sets the time.
func (f *FakeClock) SetTime(t time.Time) {
	f.lock.Lock()
	defer f.lock.Unlock()
	f.setTimeLocked(t)
}

// Actually changes the time and checks any waiters. f must be write-locked.
func (f *FakeClock) setTimeLocked(t time.Time) {
	f.time = t
	newWaiters := make([]*fakeClockWaiter, 0, len(f.waiters))
	for i := range f.waiters {
		w := f.waiters[i]
		if !w.targetTime.After(t) {

			if w.skipIfBlocked {
				select {
				case w.destChan <- t:
					w.fired = true
				default:
				}
			} else {
				w.destChan <- t
				w.fired = true
			}

			if w.stepInterval > 0 {
				for !w.targetTime.After(t) {
					w.targetTime = w.targetTime.Add(w.stepInterval)
				}
				newWaiters = append(newWaiters, w)
			}

		} else {
			newWaiters = append(newWaiters, f.waiters[i])
		}
	}
	f.waiters = newWaiters
}

// HasWaiters returns true if After has been called on f but not yet satisfied (so you can
// write race-free tests).
func (f *FakeClock) HasWaiters() bool {
	f.lock.RLock()
	defer f.lock.RUnlock()
	return len(f.waiters) > 0
}

// Sleep is akin to time.Sleep
func (f *FakeClock) Sleep(d time.Duration) {
	f.Step(d)
}

// IntervalClock implements clock.Clock, but each invocation of Now steps the clock forward the specified duration
type IntervalClock struct {
	Time     time.Time
	Duration time.Duration
}

// Now returns i's time.
func (i *IntervalClock) Now() time.Time {
	i.Time = i.Time.Add(i.Duration)
	return i.Time
}

// Since returns time since the time in i.
func (i *IntervalClock) Since(ts time.Time) time.Duration {
	return i.Time.Sub(ts)
}

// After is unimplemented, will panic.
// TODO: make interval clock use FakeClock so this can be implemented.
func (*IntervalClock) After(d time.Duration) <-chan time.Time {
	panic("IntervalClock doesn't implement After")
}

// NewTimer is unimplemented, will panic.
// TODO: make interval clock use FakeClock so this can be implemented.
func (*IntervalClock) NewTimer(d time.Duration) clock.Timer {
	panic("IntervalClock doesn't implement NewTimer")
}

// Tick is unimplemented, will panic.
// TODO: make interval clock use FakeClock so this can be implemented.
func (*IntervalClock) Tick(d time.Duration) <-chan time.Time {
	panic("IntervalClock doesn't implement Tick")
}

// Sleep is unimplemented, will panic.
func (*IntervalClock) Sleep(d time.Duration) {
	panic("IntervalClock doesn't implement Sleep")
}

var _ = clock.Timer(&fakeTimer{})

// fakeTimer implements clock.Timer based on a FakeClock.
type fakeTimer struct {
	fakeClock *FakeClock
	waiter    fakeClockWaiter
}

// C returns the channel that notifies when this timer has fired.
func (f *fakeTimer) C() <-chan time.Time {
	return f.waiter.destChan
}

// Stop stops the timer and returns true if the timer has not yet fired, or false otherwise.
func (f *fakeTimer) Stop() bool {
	f.fakeClock.lock.Lock()
	defer f.fakeClock.lock.Unlock()

	newWaiters := make([]*fakeClockWaiter, 0, len(f.fakeClock.waiters))
	for i := range f.fakeClock.waiters {
		w := f.fakeClock.waiters[i]
		if w != &f.waiter {
			newWaiters = append(newWaiters, w)
		}
	}

	f.fakeClock.waiters = newWaiters

	return !f.waiter.fired
}

// Reset resets the timer to the fake clock's "now" + d. It returns true if the timer has not yet
// fired, or false otherwise.
func (f *fakeTimer) Reset(d time.Duration) bool {
	f.fakeClock.lock.Lock()
	defer f.fakeClock.lock.Unlock()

	active := !f.waiter.fired

	f.waiter.fired = false
	f.waiter.targetTime = f.fakeClock.time.Add(d)

	var isWaiting bool
	for i := range f.fakeClock.waiters {
		w := f.fakeClock.waiters[i]
		if w == &f.waiter {
			isWaiting = true
			break
		}
	}
	if !isWaiting {
		f.fakeClock.waiters = append(f.fakeClock.waiters, &f.waiter)
	}

	return active
}