From a9de726ab29b884ce77f5375060c11d5fe6aec1d Mon Sep 17 00:00:00 2001
From: Blake Gentry <blakesgentry@gmail.com>
Date: Sat, 8 Jun 2024 13:27:20 -0500
Subject: [PATCH] expose DisableStaggerStart via TestOnlyConfig

---
 client.go      | 24 ++++++++++++++++--------
 client_test.go | 22 +++++++++++-----------
 2 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/client.go b/client.go
index 4b95cd58..8305ba01 100644
--- a/client.go
+++ b/client.go
@@ -202,6 +202,10 @@ type Config struct {
 	// Defaults to DefaultRetryPolicy.
 	RetryPolicy ClientRetryPolicy
 
+	// TestOnly holds configuration for test-only settings that should generally
+	// not be used outside of test suites.
+	TestOnly TestOnlyConfig
+
 	// Workers is a bundle of registered job workers.
 	//
 	// This field may be omitted for a program that's only enqueueing jobs
@@ -210,12 +214,6 @@ type Config struct {
 	// (i.e.  That it wasn't forgotten by accident.)
 	Workers *Workers
 
-	// Disables the normal random jittered sleep occurring in queue maintenance
-	// services to stagger their startup so they don't all try to work at the
-	// same time. Appropriate for use in tests to make sure that the client can
-	// always be started and stopped again hastily.
-	disableStaggerStart bool
-
 	// Scheduler run interval. Shared between the scheduler and producer/job
 	// executors, but not currently exposed for configuration.
 	schedulerInterval time.Duration
@@ -294,6 +292,16 @@ type QueueConfig struct {
 	MaxWorkers int
 }
 
+// TestOnlyConfig contains test-only settings that should generally not be used
+// outside of test suites.
+type TestOnlyConfig struct {
+	// DisableStaggerStart disables the normal random jittered sleep occurring in
+	// queue maintenance services to stagger their startup so they don't all try
+	// to work at the same time. Appropriate for use in tests to make sure that
+	// the client can always be started and stopped again hastily.
+	DisableStaggerStart bool
+}
+
 // Client is a single isolated instance of River. Your application may use
 // multiple instances operating on different databases or Postgres schemas
 // within a single database.
@@ -451,8 +459,8 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
 		ReindexerSchedule:           config.ReindexerSchedule,
 		RescueStuckJobsAfter:        valutil.ValOrDefault(config.RescueStuckJobsAfter, rescueAfter),
 		RetryPolicy:                 retryPolicy,
+		TestOnly:                    config.TestOnly,
 		Workers:                     config.Workers,
-		disableStaggerStart:         config.disableStaggerStart,
 		schedulerInterval:           valutil.ValOrDefault(config.schedulerInterval, maintenance.JobSchedulerIntervalDefault),
 	}
 
@@ -607,7 +615,7 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
 		// started conditionally based on whether the client is the leader.
 		client.queueMaintainer = maintenance.NewQueueMaintainer(archetype, maintenanceServices)
 
-		if config.disableStaggerStart {
+		if config.TestOnly.DisableStaggerStart {
 			client.queueMaintainer.StaggerStartupDisable(true)
 		}
 	}
diff --git a/client_test.go b/client_test.go
index 9b8cfc04..95010b8e 100644
--- a/client_test.go
+++ b/client_test.go
@@ -137,14 +137,14 @@ func newTestConfig(t *testing.T, callback callbackFunc) *Config {
 	AddWorker(workers, &noOpWorker{})
 
 	return &Config{
-		FetchCooldown:       20 * time.Millisecond,
-		FetchPollInterval:   50 * time.Millisecond,
-		Logger:              riverinternaltest.Logger(t),
-		MaxAttempts:         MaxAttemptsDefault,
-		Queues:              map[string]QueueConfig{QueueDefault: {MaxWorkers: 50}},
-		Workers:             workers,
-		disableStaggerStart: true, // disables staggered start in maintenance services
-		schedulerInterval:   riverinternaltest.SchedulerShortInterval,
+		FetchCooldown:     20 * time.Millisecond,
+		FetchPollInterval: 50 * time.Millisecond,
+		Logger:            riverinternaltest.Logger(t),
+		MaxAttempts:       MaxAttemptsDefault,
+		Queues:            map[string]QueueConfig{QueueDefault: {MaxWorkers: 50}},
+		TestOnly:          TestOnlyConfig{DisableStaggerStart: true}, // disables staggered start in maintenance services
+		Workers:           workers,
+		schedulerInterval: riverinternaltest.SchedulerShortInterval,
 	}
 }
 
@@ -3693,7 +3693,7 @@ func Test_NewClient_Defaults(t *testing.T) {
 	require.NotZero(t, client.baseService.Logger)
 	require.Equal(t, MaxAttemptsDefault, client.config.MaxAttempts)
 	require.IsType(t, &DefaultClientRetryPolicy{}, client.config.RetryPolicy)
-	require.False(t, client.config.disableStaggerStart)
+	require.False(t, client.config.TestOnly.DisableStaggerStart)
 }
 
 func Test_NewClient_Overrides(t *testing.T) {
@@ -3723,8 +3723,8 @@ func Test_NewClient_Overrides(t *testing.T) {
 		MaxAttempts:                 5,
 		Queues:                      map[string]QueueConfig{QueueDefault: {MaxWorkers: 1}},
 		RetryPolicy:                 retryPolicy,
+		TestOnly:                    TestOnlyConfig{DisableStaggerStart: true}, // disables staggered start in maintenance services
 		Workers:                     workers,
-		disableStaggerStart:         true,
 	})
 	require.NoError(t, err)
 
@@ -3745,7 +3745,7 @@ func Test_NewClient_Overrides(t *testing.T) {
 	require.Equal(t, logger, client.baseService.Logger)
 	require.Equal(t, 5, client.config.MaxAttempts)
 	require.Equal(t, retryPolicy, client.config.RetryPolicy)
-	require.True(t, client.config.disableStaggerStart)
+	require.True(t, client.config.TestOnly.DisableStaggerStart)
 }
 
 func Test_NewClient_MissingParameters(t *testing.T) {