database: GetTestDB now takes a context

database/test.go

@@ -16,5 +16,5 @@
 
 const dbLockID = 15104
 
-func clearDB(t testing.TB, c *sqlx.Conn, dsn string) {
+func clearDB(ctx context.Context, t testing.TB, c *sqlx.Conn, dsn string) {
 	// Drop all tables, and reinit
@@ -20,4 +20,4 @@
 	// Drop all tables, and reinit
-	if _, err := c.ExecContext(context.Background(), "DROP SCHEMA IF EXISTS public CASCADE"); err != nil {
+	if _, err := c.ExecContext(ctx, "DROP SCHEMA IF EXISTS public CASCADE"); err != nil {
 		t.Fatal(err)
 	}
@@ -22,6 +22,6 @@
 		t.Fatal(err)
 	}
-	if _, err := c.ExecContext(context.Background(), "CREATE SCHEMA public"); err != nil {
+	if _, err := c.ExecContext(ctx, "CREATE SCHEMA public"); err != nil {
 		t.Fatal(err)
 	}
 }
@@ -29,6 +29,7 @@
 // TestDB is a sqlx.DB wrapper that release a global lock on close
 type TestDB struct {
 	*sqlx.DB
+	ctx      context.Context
 	tb       testing.TB
 	lockConn *sqlx.Conn
 }
@@ -44,7 +45,7 @@
 
 // Clear empty the database
 func (db *TestDB) Clear() {
-	tx, err := db.lockConn.BeginTxx(context.Background(), nil)
+	tx, err := db.lockConn.BeginTxx(db.ctx, nil)
 	require.NoError(db.tb, err)
 	for _, name := range dbtablenames {
 		if _, err := tx.Exec(
@@ -72,7 +73,7 @@
 
 // GetTestDB creates a db and returns it. It must be closed within the test.
 // If it fails, t.Fatal() is called
-func GetTestDB(t testing.TB, sourceDriver source.Driver) *TestDB {
+func GetTestDB(ctx context.Context, t testing.TB, sourceDriver source.Driver) *TestDB {
 	var success bool
 
 	dsn := os.Getenv("TEST_DB_DSN")
@@ -93,7 +94,7 @@
 		}
 	}()
 
-	c, err := db.Connx(context.Background())
+	c, err := db.Connx(ctx)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -107,10 +108,10 @@
 	}()
 
 	if _, err := c.ExecContext(
-		context.Background(),
+		ctx,
 		"SELECT pg_advisory_lock($1)", dbLockID,
 	); err != nil {
 		t.Fatal(err)
 	}
 
 	if len(dbtablenames) == 0 {
@@ -111,10 +112,10 @@
 		"SELECT pg_advisory_lock($1)", dbLockID,
 	); err != nil {
 		t.Fatal(err)
 	}
 
 	if len(dbtablenames) == 0 {
-		clearDB(t, c, dsn)
+		clearDB(ctx, t, c, dsn)
 
 		m, err := NewMigrate(dsn, sourceDriver)
 		if err != nil {
@@ -163,7 +164,7 @@
 	}
 
 	success = true
-	return &TestDB{db, t, c}
+	return &TestDB{db, ctx, t, c}
 }
 
 // Close the lock connection, then the database