1//go:build integration
2// +build integration
3
4package dbtest
5
6import (
7	"fmt"
8	"os"
9	"testing"
10
11	"github.com/stretchr/testify/require"
12	"github.com/weaveworks/common/logging"
13
14	"github.com/cortexproject/cortex/pkg/configs/db"
15	"github.com/cortexproject/cortex/pkg/configs/db/postgres"
16)
17
18var (
19	done          chan error
20	dbAddr        string
21	migrationsDir string
22	errRollback   = fmt.Errorf("Rolling back test data")
23)
24
25func init() {
26	dbAddr = os.Getenv("DB_ADDR")
27	if dbAddr == "" {
28		dbAddr = "127.0.0.1"
29	}
30
31	migrationsDir = os.Getenv("MIGRATIONS_DIR")
32	if migrationsDir == "" {
33		migrationsDir = "/migrations"
34	}
35}
36
37// Setup sets up stuff for testing, creating a new database
38func Setup(t *testing.T) db.DB {
39	require.NoError(t, logging.Setup("debug"))
40	// Don't use db.MustNew, here so we can do a transaction around the whole test, to rollback.
41	pg, err := postgres.New(
42		fmt.Sprintf("postgres://postgres@%s/configs_test?sslmode=disable", dbAddr),
43		fmt.Sprintf("file:%s", migrationsDir),
44	)
45	require.NoError(t, err)
46
47	newDB := make(chan db.DB)
48	done = make(chan error)
49	go func() {
50		done <- pg.Transaction(func(tx postgres.DB) error {
51			// Pass out the tx so we can run the test
52			newDB <- tx
53			// Wait for the test to finish
54			return <-done
55		})
56	}()
57	// Get the new database
58	return <-newDB
59}
60
61// Cleanup cleans up after a test
62func Cleanup(t *testing.T, database db.DB) {
63	if done != nil {
64		done <- errRollback
65		require.Equal(t, errRollback, <-done)
66		done = nil
67	}
68	require.NoError(t, database.Close())
69}
70