165 lines
3.2 KiB
Go
165 lines
3.2 KiB
Go
package clock
|
|
|
|
import (
|
|
"errors"
|
|
"sync"
|
|
"time"
|
|
|
|
clockapi "github.com/benbjohnson/clock"
|
|
)
|
|
|
|
var (
|
|
globalClock = clockapi.New()
|
|
mu sync.Mutex
|
|
)
|
|
|
|
// Set the global clock to a clockapi.Mock with the given time.Time
|
|
func Set(t time.Time) {
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
mock, ok := globalClock.(*clockapi.Mock)
|
|
if !ok {
|
|
mock = clockapi.NewMock()
|
|
}
|
|
mock.Set(t)
|
|
globalClock = mock
|
|
}
|
|
|
|
// Add moves the mocked global clock forward the given duration. It will error
|
|
// if the global clock is not mocked.
|
|
func Add(d time.Duration) error {
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
mock, ok := globalClock.(*clockapi.Mock)
|
|
if !ok {
|
|
return errors.New("time not mocked")
|
|
}
|
|
mock.Add(d)
|
|
return nil
|
|
}
|
|
|
|
// Reset sets the global clock to a pure time implementation. Returns any
|
|
// existing Mock if set in case lingering time operations are attached to it.
|
|
func Reset() *clockapi.Mock {
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
existing := globalClock
|
|
globalClock = clockapi.New()
|
|
|
|
mock, ok := existing.(*clockapi.Mock)
|
|
if !ok {
|
|
return nil
|
|
}
|
|
return mock
|
|
}
|
|
|
|
// Clock is a non-package level wrapper around time that supports stubbing.
|
|
// It will use its localized stubs (allowing for parallelized unit tests
|
|
// where package level stubbing would cause issues). It falls back to any
|
|
// package level time stubs for non-parallel, cross-package integration
|
|
// testing scenarios.
|
|
//
|
|
// If nothing is stubbed, it defaults to default time behavior in the time
|
|
// package.
|
|
type Clock struct {
|
|
mock *clockapi.Mock
|
|
sync.Mutex
|
|
}
|
|
|
|
// Set sets the Clock to a clock.Mock at the given time.Time
|
|
func (c *Clock) Set(t time.Time) {
|
|
c.Lock()
|
|
defer c.Unlock()
|
|
if c.mock == nil {
|
|
c.mock = clockapi.NewMock()
|
|
}
|
|
c.mock.Set(t)
|
|
}
|
|
|
|
// Add moves clock forward time.Duration if it is mocked. It will error
|
|
// if the clock is not mocked.
|
|
func (c *Clock) Add(d time.Duration) error {
|
|
c.Lock()
|
|
defer c.Unlock()
|
|
if c.mock == nil {
|
|
return errors.New("clock not mocked")
|
|
}
|
|
c.mock.Add(d)
|
|
return nil
|
|
}
|
|
|
|
// Reset removes local clock.Mock. Returns any existing Mock if set in case
|
|
// lingering time operations are attached to it.
|
|
func (c *Clock) Reset() *clockapi.Mock {
|
|
c.Lock()
|
|
defer c.Unlock()
|
|
existing := c.mock
|
|
c.mock = nil
|
|
return existing
|
|
}
|
|
|
|
func (c *Clock) After(d time.Duration) <-chan time.Time {
|
|
m := c.mock
|
|
if m == nil {
|
|
return globalClock.After(d)
|
|
}
|
|
return m.After(d)
|
|
}
|
|
|
|
func (c *Clock) AfterFunc(d time.Duration, f func()) *clockapi.Timer {
|
|
m := c.mock
|
|
if m == nil {
|
|
return globalClock.AfterFunc(d, f)
|
|
}
|
|
return m.AfterFunc(d, f)
|
|
}
|
|
|
|
func (c *Clock) Now() time.Time {
|
|
m := c.mock
|
|
if m == nil {
|
|
return globalClock.Now()
|
|
}
|
|
return m.Now()
|
|
}
|
|
|
|
func (c *Clock) Since(t time.Time) time.Duration {
|
|
m := c.mock
|
|
if m == nil {
|
|
return globalClock.Since(t)
|
|
}
|
|
return m.Since(t)
|
|
}
|
|
|
|
func (c *Clock) Sleep(d time.Duration) {
|
|
m := c.mock
|
|
if m == nil {
|
|
globalClock.Sleep(d)
|
|
return
|
|
}
|
|
m.Sleep(d)
|
|
}
|
|
|
|
func (c *Clock) Tick(d time.Duration) <-chan time.Time {
|
|
m := c.mock
|
|
if m == nil {
|
|
return globalClock.Tick(d)
|
|
}
|
|
return m.Tick(d)
|
|
}
|
|
|
|
func (c *Clock) Ticker(d time.Duration) *clockapi.Ticker {
|
|
m := c.mock
|
|
if m == nil {
|
|
return globalClock.Ticker(d)
|
|
}
|
|
return m.Ticker(d)
|
|
}
|
|
|
|
func (c *Clock) Timer(d time.Duration) *clockapi.Timer {
|
|
m := c.mock
|
|
if m == nil {
|
|
return globalClock.Timer(d)
|
|
}
|
|
return m.Timer(d)
|
|
}
|