Simplify mocking of now and make sequential now func thread safe

dates/now.go

@@ -1,66 +1,51 @@
 package dates
 
 import (
+	"sync"
 	"time"
 )
 
-// Now returns the time now.. according to the current source of now
+// Now returns the time now.. according to the current now function which can be switched out for testing.
 func Now() time.Time {
-	return currentNowSource.Now()
+	return currentNow()
 }
 
 // Since returns the time elapsed since t
 func Since(t time.Time) time.Duration {
 	return Now().Sub(t)
 }
 
-// NowSource is something that can provide a now result
-type NowSource interface {
-	Now() time.Time
-}
-
-// defaultNowSource returns now as the current system time
-type defaultNowSource struct{}
-
-func (s defaultNowSource) Now() time.Time {
-	return time.Now()
-}
+// NowFunc is a function that can provide a now time
+type NowFunc func() time.Time
 
-// DefaultNowSource is the default time source
-var DefaultNowSource NowSource = defaultNowSource{}
-var currentNowSource = DefaultNowSource
+var currentNow = time.Now
 
-// SetNowSource sets the time source used by Now()
-func SetNowSource(source NowSource) {
-	currentNowSource = source
+// SetNowFunc sets the current now function
+func SetNowFunc(source NowFunc) {
+	currentNow = source
 }
 
-// a source which returns a fixed time
-type fixedNowSource struct {
-	now time.Time
+// NewFixedNow creates a new fixed now func
+func NewFixedNow(now time.Time) NowFunc {
+	return func() time.Time { return now }
 }
 
-func (s *fixedNowSource) Now() time.Time {
-	return s.now
+type sequentialNow struct {
+	start time.Time
+	step  time.Duration
+	mutex sync.Mutex
 }
 
-// NewFixedNowSource creates a new fixed time now source
-func NewFixedNowSource(now time.Time) NowSource {
-	return &fixedNowSource{now: now}
-}
-
-// a now source which returns a sequence of times 1 second after each other
-type sequentialNowSource struct {
-	current time.Time
-}
+func (s *sequentialNow) now() time.Time {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
 
-func (s *sequentialNowSource) Now() time.Time {
-	now := s.current
-	s.current = s.current.Add(time.Second * 1)
+	now := s.start
+	s.start = s.start.Add(s.step)
 	return now
 }
 
-// NewSequentialNowSource creates a new sequential time source
-func NewSequentialNowSource(start time.Time) NowSource {
-	return &sequentialNowSource{current: start}
+// NewSequentialNow creates a new sequential time func
+func NewSequentialNow(start time.Time, step time.Duration) NowFunc {
+	return (&sequentialNow{start: start, step: step}).now
 }

dates/now_test.go

@@ -9,16 +9,16 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestTimeSources(t *testing.T) {
-	defer dates.SetNowSource(dates.DefaultNowSource)
+func TestNowFuncs(t *testing.T) {
+	defer dates.SetNowFunc(time.Now)
 
 	d1 := time.Date(2018, 7, 5, 16, 29, 30, 123456, time.UTC)
-	dates.SetNowSource(dates.NewFixedNowSource(d1))
+	dates.SetNowFunc(dates.NewFixedNow(d1))
 
 	assert.Equal(t, time.Date(2018, 7, 5, 16, 29, 30, 123456, time.UTC), dates.Now())
 	assert.Equal(t, time.Date(2018, 7, 5, 16, 29, 30, 123456, time.UTC), dates.Now())
 
-	dates.SetNowSource(dates.NewSequentialNowSource(d1))
+	dates.SetNowFunc(dates.NewSequentialNow(d1, time.Second))
 
 	assert.Equal(t, time.Date(2018, 7, 5, 16, 29, 30, 123456, time.UTC), dates.Now())
 	assert.Equal(t, time.Date(2018, 7, 5, 16, 29, 31, 123456, time.UTC), dates.Now())

httpx/http_test.go

@@ -64,9 +64,9 @@ func newTestHTTPServer(port int) *httptest.Server {
 }
 
 func TestDoTrace(t *testing.T) {
-	defer dates.SetNowSource(dates.DefaultNowSource)
+	defer dates.SetNowFunc(time.Now)
 
-	dates.SetNowSource(dates.NewSequentialNowSource(time.Date(2019, 10, 7, 15, 21, 30, 123456789, time.UTC)))
+	dates.SetNowFunc(dates.NewSequentialNow(time.Date(2019, 10, 7, 15, 21, 30, 123456789, time.UTC), time.Second))
 
 	server := newTestHTTPServer(52025)
 

httpx/retries_test.go

@@ -139,9 +139,9 @@ func TestDoWithRetries(t *testing.T) {
 }
 
 func TestParseRetryAfter(t *testing.T) {
-	defer dates.SetNowSource(dates.DefaultNowSource)
+	defer dates.SetNowFunc(time.Now)
 
-	dates.SetNowSource(dates.NewFixedNowSource(time.Date(2020, 1, 7, 15, 10, 30, 500000000, time.UTC)))
+	dates.SetNowFunc(dates.NewFixedNow(time.Date(2020, 1, 7, 15, 10, 30, 500000000, time.UTC)))
 
 	assert.Equal(t, 0*time.Second, httpx.ParseRetryAfter("x"))
 	assert.Equal(t, 0*time.Second, httpx.ParseRetryAfter("0"))

uuids/uuid.go

@@ -52,12 +52,12 @@ func SetGenerator(generator Generator) {
 // generates a seedable random v4 UUID using math/rand
 type seededGenerator struct {
 	rnd *rand.Rand
-	now dates.NowSource
+	now dates.NowFunc
 }
 
 // NewSeededGenerator creates a new UUID generator that uses the given seed for the random component and the time source
 // for the time component (only applies to v7)
-func NewSeededGenerator(seed int64, now dates.NowSource) Generator {
+func NewSeededGenerator(seed int64, now dates.NowFunc) Generator {
 	return &seededGenerator{rnd: random.NewSeededGenerator(seed), now: now}
 }
 
@@ -70,7 +70,7 @@ func (g *seededGenerator) NextV4() UUID {
 func (g *seededGenerator) NextV7() UUID {
 	u := uuid.Must(uuid.NewRandomFromReader(g.rnd))
 
-	nano := g.now.Now().UnixNano()
+	nano := g.now().UnixNano()
 	t := nano / 1_000_000
 	s := (nano - t*1_000_000) >> 8
 

uuids/uuid_test.go

@@ -30,7 +30,7 @@ func TestNewV7(t *testing.T) {
 func TestSeededGenerator(t *testing.T) {
 	defer uuids.SetGenerator(uuids.DefaultGenerator)
 
-	uuids.SetGenerator(uuids.NewSeededGenerator(123456, dates.NewSequentialNowSource(time.Date(2024, 7, 32, 17, 29, 30, 123456, time.UTC))))
+	uuids.SetGenerator(uuids.NewSeededGenerator(123456, dates.NewSequentialNow(time.Date(2024, 7, 32, 17, 29, 30, 123456, time.UTC), time.Second)))
 
 	uuid1 := uuids.NewV4()
 	uuid2 := uuids.NewV7()
@@ -44,7 +44,7 @@ func TestSeededGenerator(t *testing.T) {
 	assert.Equal(t, uuids.UUID(`01910efd-5890-71e2-bd38-d266ec8d3716`), uuid2)
 	assert.Equal(t, uuids.UUID(`8720f157-ca1c-432f-9c0b-2014ddc77094`), uuid3)
 
-	uuids.SetGenerator(uuids.NewSeededGenerator(123456, dates.NewSequentialNowSource(time.Date(2024, 7, 32, 17, 29, 30, 123456, time.UTC))))
+	uuids.SetGenerator(uuids.NewSeededGenerator(123456, dates.NewSequentialNow(time.Date(2024, 7, 32, 17, 29, 30, 123456, time.UTC), time.Second)))
 
 	// should get same sequence again for same seed
 	assert.Equal(t, uuids.UUID(`d2f852ec-7b4e-457f-ae7f-f8b243c49ff5`), uuids.NewV4())