From 34be40bdcefed6c2c6054c3cef99c3ecc14808bf Mon Sep 17 00:00:00 2001
From: Jay Janssen <jayj@squareup.com>
Date: Wed, 26 Jun 2024 12:52:55 -0400
Subject: [PATCH] support optionFns to make changing refresh interval in tests
 more elegant

---
 pkg/dbconn/metadatalock.go      | 10 ++++++----
 pkg/dbconn/metadatalock_test.go |  8 ++++----
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/pkg/dbconn/metadatalock.go b/pkg/dbconn/metadatalock.go
index 5610053..933348f 100644
--- a/pkg/dbconn/metadatalock.go
+++ b/pkg/dbconn/metadatalock.go
@@ -22,7 +22,7 @@ type MetadataLock struct {
 	refreshInterval time.Duration
 }
 
-func NewMetadataLock(ctx context.Context, dsn string, lockName string, logger loggers.Advanced) (*MetadataLock, error) {
+func NewMetadataLock(ctx context.Context, dsn string, lockName string, logger loggers.Advanced, optionFns ...func(*MetadataLock)) (*MetadataLock, error) {
 	if len(lockName) == 0 {
 		return nil, errors.New("metadata lock name is empty")
 	}
@@ -34,6 +34,11 @@ func NewMetadataLock(ctx context.Context, dsn string, lockName string, logger lo
 		refreshInterval: refreshInterval,
 	}
 
+	// Apply option functions
+	for _, optionFn := range optionFns {
+		optionFn(mdl)
+	}
+
 	// Setup the dedicated connection for this lock
 	dbConfig := NewDBConfig()
 	dbConfig.MaxOpenConnections = 1
@@ -71,9 +76,6 @@ func NewMetadataLock(ctx context.Context, dsn string, lockName string, logger lo
 	ctx, mdl.cancel = context.WithCancel(ctx)
 	mdl.closeCh = make(chan error)
 	go func() {
-		// The only reason to sleep here is to give our tests a chance to override the mdl.refreshInterval
-		time.Sleep(1 * time.Second)
-
 		ticker := time.NewTicker(mdl.refreshInterval)
 		defer ticker.Stop()
 		for {
diff --git a/pkg/dbconn/metadatalock_test.go b/pkg/dbconn/metadatalock_test.go
index 28fc207..f3b6cb6 100644
--- a/pkg/dbconn/metadatalock_test.go
+++ b/pkg/dbconn/metadatalock_test.go
@@ -55,13 +55,13 @@ func TestMetadataLockContextCancel(t *testing.T) {
 func TestMetadataLockRefresh(t *testing.T) {
 	lockName := "test-refresh"
 	logger := logrus.New()
-	mdl, err := NewMetadataLock(context.Background(), testutils.DSN(), lockName, logger)
+	mdl, err := NewMetadataLock(context.Background(), testutils.DSN(), lockName, logger, func(mdl *MetadataLock) {
+		// override the refresh interval for faster testing
+		mdl.refreshInterval = 2 * time.Second
+	})
 	assert.NoError(t, err)
 	assert.NotNil(t, mdl)
 
-	// override the refresh interval to 5 seconds
-	mdl.refreshInterval = 2 * time.Second
-
 	// wait for the refresh to happen
 	time.Sleep(5 * time.Second)