use separate initializer

mafredri · mafredri · commit 60b1005f5a78 · 2023-10-12T10:19:04.000Z
diff --git a/coderd/database/dbtestutil/db.go b/coderd/database/dbtestutil/db.go
@@ -50,22 +50,34 @@ func WithDumpOnFailure() Option {
 	}
 }
 
-func WithReturnSQLDB(f func(*sql.DB)) Option {
+func withReturnSQLDB(f func(*sql.DB)) Option {
 	return func(o *options) {
 		o.returnSQLDB = f
 	}
 }
 
+func NewDBWithSQLDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub, *sql.DB) {
+	t.Helper()
+
+	if !WillUsePostgres() {
+		t.Fatalf("cannot use WithReturnSQLDB without PostgreSQL, consider adding `if !dbtestutil.WillUsePostgres() { t.Skip() }` to this test")
+	}
+
+	var sqlDB *sql.DB
+	opts = append(opts, withReturnSQLDB(func(db *sql.DB) {
+		sqlDB = db
+	}))
+	db, ps := NewDB(t, opts...)
+	return db, ps, sqlDB
+}
+
 func NewDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub) {
 	t.Helper()
 
 	var o options
 	for _, opt := range opts {
 		opt(&o)
 	}
-	if o.returnSQLDB && !WillUsePostgres() {
-		t.Fatalf("cannot use WithReturnSQLDB without PostgreSQL, consider adding `if !dbtestutil.WillUsePostgres() { t.Skip() }` to this test")
-	}
 
 	db := dbfake.New()
 	ps := pubsub.NewInMemory()
