-
Notifications
You must be signed in to change notification settings - Fork 126
/
db_testutil.go
81 lines (67 loc) · 2.46 KB
/
db_testutil.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
package database
import (
"fmt"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
_ "github.com/jackc/pgx/v5/stdlib"
"github.com/pkg/errors"
"github.com/armadaproject/armada/internal/armada/configuration"
"github.com/armadaproject/armada/internal/common/armadacontext"
"github.com/armadaproject/armada/internal/common/util"
)
// WithTestDb spins up a Postgres database for testing
// migrations: perform the list of migrations before entering the action callback
// action: callback for client code
func WithTestDb(migrations []Migration, action func(db *pgxpool.Pool) error) error {
ctx := armadacontext.Background()
// Connect and create a dedicated database for the test.
dbName := "test_" + util.NewULID()
connectionString := "host=localhost port=5432 user=postgres password=psw sslmode=disable"
db, err := pgx.Connect(ctx, connectionString)
if err != nil {
return errors.WithStack(err)
}
defer db.Close(ctx)
_, err = db.Exec(ctx, "CREATE DATABASE "+dbName)
if err != nil {
return errors.WithStack(err)
}
// Disconnect all db users and drop the database we created at test completion.
defer func() {
if _, err := db.Exec(ctx,
`SELECT pg_terminate_backend(pg_stat_activity.pid)
FROM pg_stat_activity WHERE pg_stat_activity.datname = '`+dbName+`';`,
); err != nil {
fmt.Println("failed to disconnect users:", err)
}
if _, err := db.Exec(ctx, "DROP DATABASE "+dbName); err != nil {
fmt.Println("failed to drop database:", err)
}
}()
// Connect again: this time to the database we just created. This is the database we use for tests.
testDbPool, err := pgxpool.New(ctx, connectionString+" dbname="+dbName)
if err != nil {
return errors.WithStack(err)
}
if err := UpdateDatabase(ctx, testDbPool, migrations); err != nil {
return errors.WithStack(err)
}
return action(testDbPool)
}
// WithTestDbCustom connects to specified database for testing
// migrations: perform the list of migrations before entering the action callback
// config: PostgresConfig to specify connection details to database
// action: callback for client code
func WithTestDbCustom(migrations []Migration, config configuration.PostgresConfig, action func(db *pgxpool.Pool) error) error {
ctx := armadacontext.Background()
testDbPool, err := OpenPgxPool(config)
if err != nil {
return errors.WithStack(err)
}
defer testDbPool.Close()
err = UpdateDatabase(ctx, testDbPool, migrations)
if err != nil {
return errors.WithStack(err)
}
return action(testDbPool)
}