test(coderd/database/dbtestutil): allow access to *sql.DB

mafredri · mafredri · commit 8abb2ae11634 · 2023-10-12T10:05:49.000Z
diff --git a/coderd/database/dbtestutil/db.go b/coderd/database/dbtestutil/db.go
@@ -31,6 +31,7 @@ func WillUsePostgres() bool {
 type options struct {
 	fixedTimezone string
 	dumpOnFailure bool
+	returnSQLDB   func(*sql.DB)
 }
 
 type Option func(*options)
@@ -49,6 +50,12 @@ func WithDumpOnFailure() Option {
 	}
 }
 
+func WithReturnSQLDB(f func(*sql.DB)) Option {
+	return func(o *options) {
+		o.returnSQLDB = f
+	}
+}
+
 func NewDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub) {
 	t.Helper()
 
@@ -88,6 +95,9 @@ func NewDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub) {
 		t.Cleanup(func() {
 			_ = sqlDB.Close()
 		})
+		if o.returnSQLDB != nil {
+			o.returnSQLDB(sqlDB)
+		}
 		if o.dumpOnFailure {
 			t.Cleanup(func() { DumpOnFailure(t, connectionURL) })
 		}
