expose Config.Test with Time setting

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Exposed `TestConfig` struct on `Config` under the `Test` field for configuration that is specific to test environments. For now, the only field on this type is `Time`, which can be used to set a synthetic `TimeGenerator` for tests. A stubbable time generator was added as `rivertest.TimeStub` to allow time to be easily stubbed in tests. [PR #754](https://github.com/riverqueue/river/pull/754).
+
 ### Changed
 
 - Errors returned from retryable jobs are now logged with warning logs instead of error logs. Error logs are still used for jobs that error after reaching `max_attempts`. [PR #743](https://github.com/riverqueue/river/pull/743).

client.go

@@ -46,6 +46,12 @@ const (
 	QueueNumWorkersMax = 10_000
 )
 
+// TestConfig contains configuration specific to test environments.
+type TestConfig struct {
+	// Time is a time generator to make time stubbable in tests.
+	Time rivertype.TimeGenerator
+}
+
 // Config is the configuration for a Client.
 //
 // Both Queues and Workers are required for a client to work jobs, but an
@@ -231,6 +237,9 @@ type Config struct {
 	// Defaults to false.
 	SkipUnknownJobCheck bool
 
+	// Test holds configuration specific to test environments.
+	Test TestConfig
+
 	// TestOnly can be set to true to disable certain features that are useful
 	// in production, but which may be harmful to tests, in ways like having the
 	// effect of making them slower. It should not be used outside of test
@@ -257,9 +266,6 @@ type Config struct {
 	// Scheduler run interval. Shared between the scheduler and producer/job
 	// executors, but not currently exposed for configuration.
 	schedulerInterval time.Duration
-
-	// Time generator to make time stubbable in tests.
-	time baseservice.TimeGenerator
 }
 
 func (c *Config) validate() error {
@@ -497,20 +503,24 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
 		RescueStuckJobsAfter:        valutil.ValOrDefault(config.RescueStuckJobsAfter, rescueAfter),
 		RetryPolicy:                 retryPolicy,
 		SkipUnknownJobCheck:         config.SkipUnknownJobCheck,
+		Test:                        config.Test,
 		TestOnly:                    config.TestOnly,
 		Workers:                     config.Workers,
 		WorkerMiddleware:            config.WorkerMiddleware,
 		schedulerInterval:           valutil.ValOrDefault(config.schedulerInterval, maintenance.JobSchedulerIntervalDefault),
-		time:                        config.time,
 	}
 
 	if err := config.validate(); err != nil {
 		return nil, err
 	}
 
 	archetype := baseservice.NewArchetype(config.Logger)
-	if config.time != nil {
-		archetype.Time = config.time
+	if config.Test.Time != nil {
+		if withStub, ok := config.Test.Time.(baseservice.TimeGeneratorWithStub); ok {
+			archetype.Time = withStub
+		} else {
+			archetype.Time = &baseservice.TimeGeneratorWithStubWrapper{TimeGenerator: config.Test.Time}
+		}
 	}
 
 	client := &Client[TTx]{

client_test.go

@@ -32,6 +32,7 @@ import (
 	"github.com/riverqueue/river/riverdriver"
 	"github.com/riverqueue/river/riverdriver/riverdatabasesql"
 	"github.com/riverqueue/river/riverdriver/riverpgxv5"
+	"github.com/riverqueue/river/rivershared/baseservice"
 	"github.com/riverqueue/river/rivershared/riversharedtest"
 	"github.com/riverqueue/river/rivershared/startstoptest"
 	"github.com/riverqueue/river/rivershared/testfactory"
@@ -132,10 +133,12 @@ func newTestConfig(t *testing.T, callback callbackFunc) *Config {
 		Logger:            riversharedtest.Logger(t),
 		MaxAttempts:       MaxAttemptsDefault,
 		Queues:            map[string]QueueConfig{QueueDefault: {MaxWorkers: 50}},
+		Test: TestConfig{
+			Time: &riversharedtest.TimeStub{},
+		},
 		TestOnly:          true, // disables staggered start in maintenance services
 		Workers:           workers,
 		schedulerInterval: riverinternaltest.SchedulerShortInterval,
-		time:              &riversharedtest.TimeStub{},
 	}
 }
 
@@ -5038,6 +5041,8 @@ func Test_NewClient_Defaults(t *testing.T) {
 	require.Equal(t, MaxAttemptsDefault, client.config.MaxAttempts)
 	require.IsType(t, &DefaultClientRetryPolicy{}, client.config.RetryPolicy)
 	require.False(t, client.config.SkipUnknownJobCheck)
+	require.IsType(t, nil, client.config.Test.Time)
+	require.IsType(t, &baseservice.UnStubbableTimeGenerator{}, client.baseService.Time)
 }
 
 func Test_NewClient_Overrides(t *testing.T) {

internal/dbunique/db_unique.go

@@ -9,7 +9,6 @@ import (
 	"github.com/tidwall/gjson"
 	"github.com/tidwall/sjson"
 
-	"github.com/riverqueue/river/rivershared/baseservice"
 	"github.com/riverqueue/river/rivershared/util/ptrutil"
 	"github.com/riverqueue/river/rivershared/util/sliceutil"
 	"github.com/riverqueue/river/rivertype"
@@ -54,7 +53,7 @@ func (o *UniqueOpts) StateBitmask() byte {
 	return UniqueStatesToBitmask(states)
 }
 
-func UniqueKey(timeGen baseservice.TimeGenerator, uniqueOpts *UniqueOpts, params *rivertype.JobInsertParams) ([]byte, error) {
+func UniqueKey(timeGen rivertype.TimeGenerator, uniqueOpts *UniqueOpts, params *rivertype.JobInsertParams) ([]byte, error) {
 	uniqueKeyString, err := buildUniqueKeyString(timeGen, uniqueOpts, params)
 	if err != nil {
 		return nil, err
@@ -65,7 +64,7 @@ func UniqueKey(timeGen baseservice.TimeGenerator, uniqueOpts *UniqueOpts, params
 
 // Builds a unique key made up of the unique options in place. The key is hashed
 // to become a value for `unique_key`.
-func buildUniqueKeyString(timeGen baseservice.TimeGenerator, uniqueOpts *UniqueOpts, params *rivertype.JobInsertParams) (string, error) {
+func buildUniqueKeyString(timeGen rivertype.TimeGenerator, uniqueOpts *UniqueOpts, params *rivertype.JobInsertParams) (string, error) {
 	var sb strings.Builder
 
 	if !uniqueOpts.ExcludeKind {

rivershared/baseservice/base_service.go

@@ -10,6 +10,8 @@ import (
 	"reflect"
 	"regexp"
 	"time"
+
+	"github.com/riverqueue/river/rivertype"
 )
 
 // Archetype contains the set of base service properties that are immutable, or
@@ -21,11 +23,12 @@ type Archetype struct {
 	Logger *slog.Logger
 
 	// Time returns a time generator to get the current time in UTC. Normally
-	// it's implemented as UnStubbableTimeGenerator which just calls through to
-	// `time.Now().UTC()`, but is riverinternaltest.timeStub in tests to allow
-	// the current time to be stubbed.  Services should try to use this function
-	// instead of the vanilla ones from the `time` package for testing purposes.
-	Time TimeGenerator
+	// it's implemented as [UnStubbableTimeGenerator] which just calls
+	// through to `time.Now().UTC()`, but is riverinternaltest.timeStub in tests
+	// to allow the current time to be stubbed. Services should try to use this
+	// function instead of the vanilla ones from the `time` package for testing
+	// purposes.
+	Time TimeGeneratorWithStub
 }
 
 // NewArchetype returns a new archetype. This function is most suitable for
@@ -76,32 +79,33 @@ func Init[TService withBaseService](archetype *Archetype, service TService) TSer
 	return service
 }
 
-// TimeGenerator generates a current time in UTC. In test environments it's
-// implemented by riverinternaltest.timeStub which lets the current time be
-// stubbed. Otherwise, it's implemented as UnStubbableTimeGenerator which
-// doesn't allow stubbing.
-type TimeGenerator interface {
-	// NowUTC returns the current time. This may be a stubbed time if the time
-	// has been actively stubbed in a test.
-	NowUTC() time.Time
-
-	// NowUTCOrNil returns if the currently stubbed time _if_ the current time
-	// is stubbed, and returns nil otherwise. This is generally useful in cases
-	// where a component may want to use a stubbed time if the time is stubbed,
-	// but to fall back to a database time default otherwise.
-	NowUTCOrNil() *time.Time
+type TimeGeneratorWithStub interface {
+	rivertype.TimeGenerator
 
 	// StubNowUTC stubs the current time. It will panic if invoked outside of
 	// tests. Returns the same time passed as parameter for convenience.
 	StubNowUTC(nowUTC time.Time) time.Time
 }
 
+// TimeGeneratorWithStubWrapper provides a wrapper around TimeGenerator that
+// implements missing TimeGeneratorWithStub functions. This is used so that we
+// only need to expose the minimal TimeGenerator interface publicly, but can
+// keep a stubbable version of widely available for internal use.
+type TimeGeneratorWithStubWrapper struct {
+	rivertype.TimeGenerator
+}
+
+func (g *TimeGeneratorWithStubWrapper) StubNowUTC(nowUTC time.Time) time.Time {
+	panic("time not stubbable outside tests")
+}
+
 // UnStubbableTimeGenerator is a TimeGenerator implementation that can't be
 // stubbed. It's always the generator used outside of tests.
 type UnStubbableTimeGenerator struct{}
 
 func (g *UnStubbableTimeGenerator) NowUTC() time.Time       { return time.Now() }
 func (g *UnStubbableTimeGenerator) NowUTCOrNil() *time.Time { return nil }
+
 func (g *UnStubbableTimeGenerator) StubNowUTC(nowUTC time.Time) time.Time {
 	panic("time not stubbable outside tests")
 }

rivershared/riversharedtest/riversharedtest.go

@@ -48,8 +48,10 @@ func LoggerWarn(tb testing.TB) *slog.Logger {
 	return slogtest.NewLogger(tb, &slog.HandlerOptions{Level: slog.LevelWarn})
 }
 
-// TimeStub implements baseservice.TimeGenerator to allow time to be stubbed in
-// tests.
+// TimeStub implements baseservice.TimeGeneratorWithStub to allow time to be
+// stubbed in tests.
+//
+// It exists separately from rivertest.TimeStub to avoid a circular dependency.
 type TimeStub struct {
 	mu     sync.RWMutex
 	nowUTC *time.Time

rivertest/time_stub.go

@@ -0,0 +1,48 @@
+package rivertest
+
+import (
+	"sync"
+	"time"
+)
+
+// TimeStub implements rivertype.TimeGenerator to allow time to be stubbed in
+// tests. It is implemented in a thread-safe manner with a mutex, allowing the
+// current time to be stubbed at any time with StubNowUTC.
+type TimeStub struct {
+	mu     sync.RWMutex
+	nowUTC *time.Time
+}
+
+// NowUTC returns the current time. This may be a stubbed time if the time has
+// been actively stubbed in a test.
+func (t *TimeStub) NowUTC() time.Time {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+
+	if t.nowUTC == nil {
+		return time.Now().UTC()
+	}
+
+	return *t.nowUTC
+}
+
+// NowUTCOrNil returns if the currently stubbed time _if_ the current time
+// is stubbed, and returns nil otherwise. This is generally useful in cases
+// where a component may want to use a stubbed time if the time is stubbed,
+// but to fall back to a database time default otherwise.
+func (t *TimeStub) NowUTCOrNil() *time.Time {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+
+	return t.nowUTC
+}
+
+// StubNowUTC stubs the current time. It will panic if invoked outside of tests.
+// Returns the same time passed as parameter for convenience.
+func (t *TimeStub) StubNowUTC(nowUTC time.Time) time.Time {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	t.nowUTC = &nowUTC
+	return nowUTC
+}

rivertest/time_stub_test.go

@@ -0,0 +1,27 @@
+package rivertest
+
+import (
+	"testing"
+	"time"
+
+	"github.com/riverqueue/river/rivertype"
+	"github.com/stretchr/testify/require"
+)
+
+// Ensure that TimeStub implements rivertype.TimeGenerator.
+var _ rivertype.TimeGenerator = &TimeStub{}
+
+func TestTimeStub(t *testing.T) {
+	t.Parallel()
+
+	stub := &TimeStub{}
+
+	now := time.Now().UTC()
+
+	require.WithinDuration(t, now, stub.NowUTC(), 2*time.Second)
+	require.Nil(t, stub.NowUTCOrNil())
+
+	stub.StubNowUTC(now)
+	require.Equal(t, now, stub.NowUTC())
+	require.Equal(t, &now, stub.NowUTCOrNil())
+}

rivertype/time_generator.go

@@ -0,0 +1,19 @@
+package rivertype
+
+import "time"
+
+// TimeGenerator generates a current time in UTC. In test environments it's
+// implemented by riverinternaltest.timeStub which lets the current time be
+// stubbed. Otherwise, it's implemented as UnStubbableTimeGenerator which
+// doesn't allow stubbing.
+type TimeGenerator interface {
+	// NowUTC returns the current time. This may be a stubbed time if the time
+	// has been actively stubbed in a test.
+	NowUTC() time.Time
+
+	// NowUTCOrNil returns if the currently stubbed time _if_ the current time
+	// is stubbed, and returns nil otherwise. This is generally useful in cases
+	// where a component may want to use a stubbed time if the time is stubbed,
+	// but to fall back to a database time default otherwise.
+	NowUTCOrNil() *time.Time
+}