From 8abb2ae11634ab38c972d5c97967ead824033611 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Thu, 12 Oct 2023 10:05:49 +0000 Subject: [PATCH 1/3] test(coderd/database/dbtestutil): allow access to *sql.DB --- coderd/database/dbtestutil/db.go | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/coderd/database/dbtestutil/db.go b/coderd/database/dbtestutil/db.go index a32d7b31245f0..eabaedc55e5c4 100644 --- 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) }) } From 866d3d9e07582988bd1d6d61a1ca0d7ac1b20b71 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Thu, 12 Oct 2023 10:13:57 +0000 Subject: [PATCH 2/3] fatal if returnsqldb and no postgres --- coderd/database/dbtestutil/db.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/coderd/database/dbtestutil/db.go b/coderd/database/dbtestutil/db.go index eabaedc55e5c4..aa0e37d80b259 100644 --- a/coderd/database/dbtestutil/db.go +++ b/coderd/database/dbtestutil/db.go @@ -63,6 +63,9 @@ func NewDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub) { 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() From ee86b3a1bbe9128146cfa65ef74ab2f0ffd16aee Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Thu, 12 Oct 2023 10:19:04 +0000 Subject: [PATCH 3/3] use separate initializer --- coderd/database/dbtestutil/db.go | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/coderd/database/dbtestutil/db.go b/coderd/database/dbtestutil/db.go index aa0e37d80b259..65e1afecc7948 100644 --- a/coderd/database/dbtestutil/db.go +++ b/coderd/database/dbtestutil/db.go @@ -50,12 +50,27 @@ 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.Fatal("cannot use NewDBWithSQLDB 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() @@ -63,9 +78,6 @@ func NewDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub) { 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()