104 lines
2.3 KiB
Go
104 lines
2.3 KiB
Go
package timetools
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// TimeProvider is an interface we use to mock time in tests.
|
|
type TimeProvider interface {
|
|
UtcNow() time.Time
|
|
Sleep(time.Duration)
|
|
After(time.Duration) <-chan time.Time
|
|
}
|
|
|
|
// RealTime is a real clock time, used in production.
|
|
type RealTime struct {
|
|
}
|
|
|
|
func (*RealTime) UtcNow() time.Time {
|
|
return time.Now().UTC()
|
|
}
|
|
|
|
func (*RealTime) Sleep(d time.Duration) {
|
|
time.Sleep(d)
|
|
}
|
|
|
|
func (*RealTime) After(d time.Duration) <-chan time.Time {
|
|
return time.After(d)
|
|
}
|
|
|
|
// FreezedTime is manually controlled time for use in tests.
|
|
type FreezedTime struct {
|
|
CurrentTime time.Time
|
|
}
|
|
|
|
func (t *FreezedTime) UtcNow() time.Time {
|
|
return t.CurrentTime
|
|
}
|
|
|
|
func (t *FreezedTime) Sleep(d time.Duration) {
|
|
t.CurrentTime = t.CurrentTime.Add(d)
|
|
}
|
|
|
|
func (t *FreezedTime) After(d time.Duration) <-chan time.Time {
|
|
t.Sleep(d)
|
|
c := make(chan time.Time, 1)
|
|
c <- t.CurrentTime
|
|
return c
|
|
}
|
|
|
|
type sleepableTime struct {
|
|
currentTime time.Time
|
|
waiters map[time.Time][]chan time.Time
|
|
mu sync.Mutex
|
|
}
|
|
|
|
// SleepProvider returns a TimeProvider that has good fakes for
|
|
// time.Sleep and time.After. Both functions will behave as if
|
|
// time is frozen until you call AdvanceTimeBy, at which point
|
|
// any calls to time.Sleep that should return do return and
|
|
// any ticks from time.After that should happen do happen.
|
|
func SleepProvider(currentTime time.Time) TimeProvider {
|
|
return &sleepableTime{
|
|
currentTime: currentTime,
|
|
waiters: make(map[time.Time][]chan time.Time),
|
|
}
|
|
}
|
|
|
|
func (t *sleepableTime) UtcNow() time.Time {
|
|
return t.currentTime
|
|
}
|
|
|
|
func (t *sleepableTime) Sleep(d time.Duration) {
|
|
<-t.After(d)
|
|
}
|
|
|
|
func (t *sleepableTime) After(d time.Duration) <-chan time.Time {
|
|
t.mu.Lock()
|
|
defer t.mu.Unlock()
|
|
|
|
c := make(chan time.Time, 1)
|
|
until := t.currentTime.Add(d)
|
|
t.waiters[until] = append(t.waiters[until], c)
|
|
return c
|
|
}
|
|
|
|
// AdvanceTimeBy simulates advancing time by some time.Duration d.
|
|
// This function panics if st is not the result of a call to
|
|
// SleepProvider.
|
|
func AdvanceTimeBy(st TimeProvider, d time.Duration) {
|
|
t := st.(*sleepableTime)
|
|
t.mu.Lock()
|
|
defer t.mu.Unlock()
|
|
|
|
t.currentTime = t.currentTime.Add(d)
|
|
for k, v := range t.waiters {
|
|
if k.Before(t.currentTime) {
|
|
for _, c := range v {
|
|
c <- t.currentTime
|
|
}
|
|
delete(t.waiters, k)
|
|
}
|
|
}
|
|
}
|