Repository: google/exposure-notifications-server
Pull request #428 (2020-05-24 20:29:43)

Fix database tests
Database tests haven't been running since https://github.com/google/exposure-notifications-server/commit/f113c0d2#diff-f7ac49848dfd0ce5e8f3251a560f5a0f, because of the package being renamed to no longer include a _test.go suffix. Not shown in that diff is the inclusion of a TestMain, which only runs if the file is actually a test (ends in the _test.go suffix). This and other dependent database tests have been silently skipping.

This refactors the tests to use a new Docker container per instance (where test implementers decide how they want to manage the instance), and updates all the database tests to run in parallel. This results in a net increase of 8s to the overall test suite runtime on my local machine, but seems worth it given this is the 4th (5th?) time a seemingly unrelated change as silently broken the database tests.

Folks can opt of out running the tests (either using -short or setting an environment variable), but the default behavior is now to run the tests. The only requirement is Docker (which was already a requirement anyway).

google-prow-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/google/exposure-notifications-server/pull/428#" title="Author self-approved">sethvargo</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=google%2Fexposure-notifications-server).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/google/exposure-notifications-server/blob/master/OWNERS)~~ [sethvargo]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
sethvargo: Wow, barely any additional timing increase on CI either 😄 
google-prow-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/google/exposure-notifications-server/pull/428#pullrequestreview-417390410" title="LGTM">mikehelmick</a>*, *<a href="https://github.com/google/exposure-notifications-server/pull/428#" title="Author self-approved">sethvargo</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=google%2Fexposure-notifications-server).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/google/exposure-notifications-server/blob/master/OWNERS)~~ [mikehelmick,sethvargo]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->

Patch:
diff --git a/go.mod b/go.mod
index fb3f04b..52b37bb 100644
--- a/go.mod
+++ b/go.mod
@@ -44,6 +44,7 @@ require (
 	github.com/mitchellh/mapstructure v1.3.1 // indirect
 	github.com/oklog/run v1.1.0 // indirect
 	github.com/oracle/oci-go-sdk v19.3.0+incompatible // indirect
+	github.com/ory/dockertest v3.3.5+incompatible
 	github.com/pierrec/lz4 v2.5.2+incompatible // indirect
 	github.com/posener/complete v1.2.3 // indirect
 	github.com/prometheus/client_golang v1.6.0 // indirect
diff --git a/internal/database/database_util.go b/internal/database/database_util.go
new file mode 100644
index 0000000..b01e2e5
--- /dev/null
+++ b/internal/database/database_util.go
@@ -0,0 +1,162 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package database
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"net/url"
+	"os"
+	"path/filepath"
+	"runtime"
+	"strconv"
+	"testing"
+
+	"github.com/golang-migrate/migrate/v4"
+	"github.com/jackc/pgx/v4/pgxpool"
+	"github.com/ory/dockertest"
+
+	// imported to register the postgres migration driver
+	_ "github.com/golang-migrate/migrate/v4/database/postgres"
+	// imported to register the "file" source migration driver
+	_ "github.com/golang-migrate/migrate/v4/source/file"
+	// imported to register the "postgres" database driver for migrate
+)
+
+// NewTestDatabase creates a new database suitable for use in testing. This
+// should not be used outside of testing, but it is exposed in the main package
+// so it can be shared with other packages.
+//
+// All database tests can be skipped by running `go test -short` or by setting
+// the `SKIP_DATABASE_TESTS` environment variable.
+func NewTestDatabase(tb testing.TB) *DB {
+	tb.Helper()
+
+	if testing.Short() {
+		tb.Skipf("🚧 Skipping database tests (short!")
+	}
+
+	if skip, _ := strconv.ParseBool(os.Getenv("SKIP_DATABASE_TESTS")); skip {
+		tb.Skipf("🚧 Skipping database tests (SKIP_DATABASE_TESTS is set)!")
+	}
+
+	// Create the pool (docker instance).
+	pool, err := dockertest.NewPool("")
+	if err != nil {
+		tb.Fatalf("failed to create Docker pool: %s", err)
+	}
+
+	// Start the container.
+	dbname, username, password := "en-server", "my-username", "abcd1234"
+	container, err := pool.RunWithOptions(&dockertest.RunOptions{
+		Repository: "postgres",
+		Tag:        "12-alpine",
+		Env: []string{
+			"LANG=C",
+			"POSTGRES_DB=" + dbname,
+			"POSTGRES_USER=" + username,
+			"POSTGRES_PASSWORD=" + password,
+		},
+	})
+	if err != nil {
+		tb.Fatalf("failed to start postgres container: %s", err)
+	}
+
+	// Ensure container is cleaned up.
+	tb.Cleanup(func() {
+		if err := pool.Purge(container); err != nil {
+			tb.Fatalf("failed to cleanup postgres container: %s", err)
+		}
+	})
+
+	// Get the host. On Mac, Docker runs in a VM.
+	host := container.Container.NetworkSettings.IPAddress
+	if runtime.GOOS == "darwin" {
+		host = net.JoinHostPort(container.GetBoundIP("5432/tcp"), container.GetPort("5432/tcp"))
+	}
+
+	// Build the connection URL.
+	connURL := &url.URL{
+		Scheme: "postgres",
+		User:   url.UserPassword(username, password),
+		Host:   host,
+		Path:   dbname,
+	}
+	q := connURL.Query()
+	q.Add("sslmode", "disable")
+	connURL.RawQuery = q.Encode()
+
+	// Establish a connection to the database.
+	var dbpool *pgxpool.Pool
+	if err := pool.Retry(func() error {
+		var err error
+		dbpool, err = pgxpool.Connect(context.Background(), connURL.String())
+		if err != nil {
+			return fmt.Errorf("failed to connect to database: %w", err)
+		}
+		return nil
+	}); err != nil {
+		tb.Fatalf("failed to start postgres: %s", err)
+	}
+
+	// Run the migrations.
+	if err := dbMigrate(connURL.String()); err != nil {
+		tb.Fatalf("failed to migrate databse: %s", err)
+	}
+
+	// Create the db instance.
+	db := &DB{Pool: dbpool}
+
+	// Close db when done.
+	tb.Cleanup(func() {
+		db.Close(context.Background())
+	})
+
+	return db
+}
+
+// dbMigrate runs the migrations. u is the connection URL string (e.g.
+// postgres://...).
+func dbMigrate(u string) error {
+	// Run the migrations
+	migrationsDir := fmt.Sprintf("file://%s", dbMigrationsDir())
+	m, err := migrate.New(migrationsDir, u)
+	if err != nil {
+		return fmt.Errorf("failed create migrate: %w", err)
+	}
+	if err := m.Up(); err != nil && err != migrate.ErrNoChange {
+		return fmt.Errorf("failed run migrate: %w", err)
+	}
+	srcErr, dbErr := m.Close()
+	if srcErr != nil {
+		return fmt.Errorf("migrate source error: %w", srcErr)
+	}
+	if dbErr != nil {
+		return fmt.Errorf("migrate database error: %w", dbErr)
+	}
+	return nil
+}
+
+// dbMigrationsDir returns the path on disk to the migrations. It uses
+// runtime.Caller() to get the path to the caller, since this package is
+// imported by multiple others at different levels.
+func dbMigrationsDir() string {
+	_, filename, _, ok := runtime.Caller(1)
+	if !ok {
+		return ""
+	}
+	return filepath.Join(filepath.Dir(filename), "../../migrations")
+}
diff --git a/scripts/presubmit.sh b/scripts/presubmit.sh
index acded03..504b98d 100755
--- a/scripts/presubmit.sh
+++ b/scripts/presubmit.sh
@@ -101,21 +101,11 @@ echo "🚧 Compile"
 go build ./...
 
 
-echo "📚 Starting database"
-export DB_NAME="en-server-db-test"
-export DB_PORT=5435
-${ROOT}/scripts/dev dbstart && sleep 2
-${ROOT}/scripts/dev dbmigrate
-trap "${ROOT}/scripts/dev dbstop" EXIT
-
-if [ "${DB_USER:-}" == "" ]; then
-   echo "🚨 DB_USER is not configured. Test DB tests will not run."
-fi
-
-
 echo "🧪 Test"
 go test ./... \
   -coverprofile=coverage.out \
+  -count=1 \
+  -parallel=20 \
   -timeout=5m
 
 


Test patch:
diff --git a/internal/authorizedapp/database/authorized_app_test.go b/internal/authorizedapp/database/authorized_app_test.go
index 5ad30dd..344ff43 100644
--- a/internal/authorizedapp/database/authorized_app_test.go
+++ b/internal/authorizedapp/database/authorized_app_test.go
@@ -22,13 +22,11 @@ import (
 	"crypto/x509"
 	"encoding/pem"
 	"fmt"
-	"log"
-	"os"
 	"testing"
 	"time"
 
 	"github.com/google/exposure-notifications-server/internal/authorizedapp/model"
-	coredb "github.com/google/exposure-notifications-server/internal/database"
+	"github.com/google/exposure-notifications-server/internal/database"
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 )
@@ -45,26 +43,10 @@ func (s *testSecretManager) GetSecretValue(ctx context.Context, name string) (st
 	return v, nil
 }
 
-var testDB *coredb.DB
-
-func TestMain(m *testing.M) {
-	ctx := context.Background()
-
-	if os.Getenv("DB_USER") != "" {
-		var err error
-		testDB, err = coredb.CreateTestDB(ctx, "authorizedapp")
-		if err != nil {
-			log.Fatalf("creating test DB: %v", err)
-		}
-	}
-	os.Exit(m.Run())
-}
-
 func TestAuthorizedAppLifecycle(t *testing.T) {
-	if testDB == nil {
-		t.Skip("no test DB")
-	}
-	defer coredb.ResetTestDB(t, testDB)
+	t.Parallel()
+
+	testDB := database.NewTestDatabase(t)
 	ctx := context.Background()
 	aadb := New(testDB)
 	sm := &testSecretManager{
@@ -105,10 +87,8 @@ func TestAuthorizedAppLifecycle(t *testing.T) {
 }
 
 func TestGetAuthorizedApp(t *testing.T) {
-	if testDB == nil {
-		t.Skip("no test DB")
-	}
-	defer coredb.ResetTestDB(t, testDB)
+	t.Parallel()
+
 	ctx := context.Background()
 
 	// Create private key for parsing later
@@ -271,13 +251,16 @@ func TestGetAuthorizedApp(t *testing.T) {
 		c := c
 
 		t.Run(c.name, func(t *testing.T) {
+			t.Parallel()
+
+			testDB := database.NewTestDatabase(t)
+
 			// Acquire a connection
 			conn, err := testDB.Pool.Acquire(ctx)
 			if err != nil {
 				t.Fatal(err)
 			}
 			defer conn.Release()
-			defer coredb.ResetTestDB(t, testDB)
 
 			// Insert the data
 			if _, err := conn.Exec(ctx, c.sql, c.args...); err != nil {
diff --git a/internal/database/database_test_util.go b/internal/database/database_test_util.go
deleted file mode 100644
index 8c7d35f..0000000
--- a/internal/database/database_test_util.go
+++ /dev/null
@@ -1,130 +0,0 @@
-// Copyright 2020 Google LLC
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package database
-
-import (
-	"context"
-	"fmt"
-	"log"
-	"os"
-	"testing"
-
-	"github.com/golang-migrate/migrate/v4"
-	"github.com/kelseyhightower/envconfig"
-
-	// imported to register the postgres migration driver
-	_ "github.com/golang-migrate/migrate/v4/database/postgres"
-	// imported to register the "file" source migration driver
-	_ "github.com/golang-migrate/migrate/v4/source/file"
-	// imported to register the "postgres" database driver for migrate
-)
-
-var testDB *DB
-
-func TestMain(m *testing.M) {
-	ctx := context.Background()
-
-	if os.Getenv("DB_USER") != "" {
-		var err error
-		testDB, err = CreateTestDB(ctx, "core")
-		if err != nil {
-			log.Fatalf("creating test DB: %v", err)
-		}
-	}
-	os.Exit(m.Run())
-}
-
-// openTestDB connects to the Postgres server specified by the DB_XXX environment
-// variables, creates an empty test database on it, and returns a *DB connected
-// to that database.
-func CreateTestDB(ctx context.Context, suffix string) (*DB, error) {
-	var testDBName = fmt.Sprintf("exposure-server-test-%s", suffix)
-
-	// Connect to the default database to create the test database.
-	var config Config
-	err := envconfig.Process("dbtestEnvironment", &config)
-	if err != nil {
-		return nil, err
-	}
-	// Override DB name.
-	config.Name = "postgres"
-	db, err := NewFromEnv(ctx, &config)
-	if err != nil {
-		return nil, err
-	}
-	if err := createDatabase(ctx, db, testDBName); err != nil {
-		return nil, err
-	}
-	db.Close(ctx)
-
-	// Connect to the test database and create its schema by applying all migrations.
-	config.Name = testDBName
-	db, err = NewFromEnv(ctx, &config)
-	if err != nil {
-		return nil, err
-	}
-	const source = "file://../../../migrations"
-	uri := DbURI(&config)
-	m, err := migrate.New(source, uri)
-	if err != nil {
-		return nil, err
-	}
-	if err := m.Up(); err != nil && err != migrate.ErrNoChange {
-		_, _ = m.Close()
-		return nil, err
-	}
-	srcErr, dbErr := m.Close()
-	if srcErr != nil {
-		return nil, srcErr
-	}
-	if dbErr != nil {
-		return nil, dbErr
-	}
-	return db, nil
-}
-
-func createDatabase(ctx context.Context, db *DB, name string) error {
-	conn, err := db.Pool.Acquire(ctx)
-	if err != nil {
-		return err
-	}
-	defer conn.Release()
-
-	if _, err := conn.Exec(ctx, fmt.Sprintf(`DROP DATABASE IF EXISTS %q`, name)); err != nil {
-		return err
-	}
-	_, err = conn.Exec(ctx, fmt.Sprintf(`CREATE DATABASE %q`, name))
-	return err
-}
-
-func ResetTestDB(t *testing.T, testDB *DB) {
-	t.Helper()
-	ctx := context.Background()
-	conn, err := testDB.Pool.Acquire(ctx)
-	if err != nil {
-		t.Fatal(err)
-	}
-	defer conn.Release()
-
-	_, err = conn.Exec(ctx, `
-		TRUNCATE
-			FederationInQuery, FederationInSync, FederationOutAuthorization,
-			Exposure, AuthorizedApp,
-			ExportConfig, ExportBatch, ExportFile
-	`)
-	if err != nil {
-		t.Fatal(err)
-	}
-}
diff --git a/internal/database/lock_test.go b/internal/database/lock_test.go
index 6348288..d9da0ea 100644
--- a/internal/database/lock_test.go
+++ b/internal/database/lock_test.go
@@ -22,9 +22,9 @@ import (
 )
 
 func TestLock(t *testing.T) {
-	if testDB == nil {
-		t.Skip("no test DB")
-	}
+	t.Parallel()
+
+	testDB := NewTestDatabase(t)
 	ctx := context.Background()
 
 	const (
diff --git a/internal/export/database/export_test.go b/internal/export/database/export_test.go
index 28517b8..b5a223a 100644
--- a/internal/export/database/export_test.go
+++ b/internal/export/database/export_test.go
@@ -17,8 +17,6 @@ package database
 import (
 	"context"
 	"errors"
-	"log"
-	"os"
 	"sort"
 	"testing"
 	"time"
@@ -31,26 +29,10 @@ import (
 	pgx "github.com/jackc/pgx/v4"
 )
 
-var testDB *database.DB
-
-func TestMain(m *testing.M) {
-	ctx := context.Background()
-
-	if os.Getenv("DB_USER") != "" {
-		var err error
-		testDB, err = database.CreateTestDB(ctx, "export")
-		if err != nil {
-			log.Fatalf("creating test DB: %v", err)
-		}
-	}
-	os.Exit(m.Run())
-}
-
 func TestAddSignatureInfo(t *testing.T) {
-	if testDB == nil {
-		t.Skip("no test DB")
-	}
-	defer database.ResetTestDB(t, testDB)
+	t.Parallel()
+
+	testDB := database.NewTestDatabase(t)
 	ctx := context.Background()
 
 	thruTime := time.Now().UTC().Add(6 * time.Hour).Truncate(time.Microsecond)
@@ -87,10 +69,9 @@ func TestAddSignatureInfo(t *testing.T) {
 }
 
 func TestLookupSignatureInfos(t *testing.T) {
-	if testDB == nil {
-		t.Skip("no test DB")
-	}
-	defer database.ResetTestDB(t, testDB)
+	t.Parallel()
+
+	testDB := database.NewTestDatabase(t)
 	ctx := context.Background()
 
 	testTime := time.Now().UTC()
@@ -132,10 +113,9 @@ func TestLookupSignatureInfos(t *testing.T) {
 }
 
 func TestAddExportConfig(t *testing.T) {
-	if testDB == nil {
-		t.Skip("no test DB")
-	}
-	defer database.ResetTestDB(t, testDB)
+	t.Parallel()
+
+	testDB := database.NewTestDatabase(t)
 	ctx := context.Background()
 
 	fromTime := time.Now()
@@ -182,10 +162,9 @@ func TestAddExportConfig(t *testing.T) {
 }
 
 func TestIterateExportConfigs(t *testing.T) {
-	if testDB == nil {
-		t.Skip("no test DB")
-	}
-	defer database.ResetTestDB(t, testDB)
+	t.Parallel()
+
+	testDB := database.NewTestDatabase(t)
 	ctx := context.Background()
 
 	now := time.Now().Truncate(time.Microsecond)
@@ -238,10 +217,9 @@ func TestIterateExportConfigs(t *testing.T) {
 }
 
 func TestBatches(t *testing.T) {
-	if testDB == nil {
-		t.Skip("no test DB")
-	}
-	defer database.ResetTestDB(t, testDB)
+	t.Parallel()
+
+	testDB := database.NewTestDatabase(t)
 	ctx := context.Background()
 
 	now := time.Now().Truncate(time.Microsecond)
@@ -348,11 +326,10 @@ func TestBatches(t *testing.T) {
 }
 
 func TestFinalizeBatch(t *testing.T) {
-	if testDB == nil {
-		t.Skip("no test DB")
-	}
+	t.Parallel()
+
+	testDB := database.NewTestDatabase(t)
 	exportDB := New(testDB)
-	defer database.ResetTestDB(t, testDB)
 	ctx := context.Background()
 	now := time.Now().Truncate(time.Microsecond)
 
@@ -443,10 +420,9 @@ func TestFinalizeBatch(t *testing.T) {
 
 // TestKeysInBatch ensures that keys are fetched in the correct batch when they fall on boundary conditions.
 func TestKeysInBatch(t *testing.T) {
-	if testDB == nil {
-		t.Skip("no test DB")
-	}
-	defer database.ResetTestDB(t, testDB)
+	t.Parallel()
+
+	testDB := database.NewTestDatabase(t)
 	ctx := context.Background()
 	now := time.Now()
 
@@ -540,10 +516,9 @@ func TestKeysInBatch(t *testing.T) {
 
 // TestAddExportFileSkipsDuplicates ensures that ExportFile records are not overwritten.
 func TestAddExportFileSkipsDuplicates(t *testing.T) {
-	if testDB == nil {
-		t.Skip("no test DB")
-	}
-	defer database.ResetTestDB(t, testDB)
+	t.Parallel()
+
+	testDB := database.NewTestDatabase(t)
 	exportDB := New(testDB)
 	ctx := context.Background()
 
diff --git a/internal/federationin/database/federationin_test.go b/internal/federationin/database/federationin_test.go
index 05f8fde..b574534 100644
--- a/internal/federationin/database/federationin_test.go
+++ b/internal/federationin/database/federationin_test.go
@@ -17,11 +17,10 @@ package database
 import (
 	"context"
 	"errors"
-	"log"
-	"os"
 	"testing"
 	"time"
 
+	"github.com/google/exposure-notifications-server/internal/database"
 	coredb "github.com/google/exposure-notifications-server/internal/database"
 	"github.com/google/exposure-notifications-server/internal/federationin/model"
 
@@ -29,27 +28,11 @@ import (
 	"github.com/google/go-cmp/cmp/cmpopts"
 )
 
-var testDB *coredb.DB
-
-func TestMain(m *testing.M) {
-	ctx := context.Background()
-
-	if os.Getenv("DB_USER") != "" {
-		var err error
-		testDB, err = coredb.CreateTestDB(ctx, "federationin")
-		if err != nil {
-			log.Fatalf("creating test DB: %v", err)
-		}
-	}
-	os.Exit(m.Run())
-}
-
 // TestFederationIn tests functions operating over FederationInQuery, FederationInSync.
 func TestFederationIn(t *testing.T) {
-	if testDB == nil {
-		t.Skip("no test DB")
-	}
-	defer coredb.ResetTestDB(t, testDB)
+	t.Parallel()
+
+	testDB := database.NewTestDatabase(t)
 	ctx := context.Background()
 
 	ts := time.Date(2020, 5, 6, 0, 0, 0, 0, time.UTC)
diff --git a/internal/federationout/database/federationout_test.go b/internal/federationout/database/federationout_test.go
index a9533d5..9024d69 100644
--- a/internal/federationout/database/federationout_test.go
+++ b/internal/federationout/database/federationout_test.go
@@ -17,37 +17,20 @@ package database
 import (
 	"context"
 	"errors"
-	"log"
-	"os"
 	"testing"
 
+	"github.com/google/exposure-notifications-server/internal/database"
 	coredb "github.com/google/exposure-notifications-server/internal/database"
 	"github.com/google/exposure-notifications-server/internal/federationin/model"
 
 	"github.com/google/go-cmp/cmp"
 )
 
-var testDB *coredb.DB
-
-func TestMain(m *testing.M) {
-	ctx := context.Background()
-
-	if os.Getenv("DB_USER") != "" {
-		var err error
-		testDB, err = coredb.CreateTestDB(ctx, "federationout")
-		if err != nil {
-			log.Fatalf("creating test DB: %v", err)
-		}
-	}
-	os.Exit(m.Run())
-}
-
 // TestFederationOutAuthorization tests the functions accessing the FederationOutAuthorization table.
 func TestFederationOutAuthorization(t *testing.T) {
-	if testDB == nil {
-		t.Skip("no test DB")
-	}
-	defer coredb.ResetTestDB(t, testDB)
+	t.Parallel()
+
+	testDB := database.NewTestDatabase(t)
 	ctx := context.Background()
 
 	want := &model.FederationOutAuthorization{
diff --git a/internal/publish/database/exposure_test.go b/internal/publish/database/exposure_test.go
index 0b039bf..b2f5335 100644
--- a/internal/publish/database/exposure_test.go
+++ b/internal/publish/database/exposure_test.go
@@ -17,38 +17,21 @@ package database
 import (
 	"context"
 	"errors"
-	"log"
-	"os"
+	"fmt"
 	"testing"
 	"time"
 
-	coredb "github.com/google/exposure-notifications-server/internal/database"
+	"github.com/google/exposure-notifications-server/internal/database"
 
 	"github.com/google/exposure-notifications-server/internal/publish/model"
 	"github.com/google/go-cmp/cmp"
 )
 
-var testDB *coredb.DB
-var testPublishDB *PublishDB
-
-func TestMain(m *testing.M) {
-	ctx := context.Background()
-
-	if os.Getenv("DB_USER") != "" {
-		var err error
-		testDB, err = coredb.CreateTestDB(ctx, "publish")
-		if err != nil {
-			log.Fatalf("creating test DB: %v", err)
-		}
-		testPublishDB = New(testDB)
-	}
-	os.Exit(m.Run())
-}
 func TestExposures(t *testing.T) {
-	if testDB == nil {
-		t.Skip("no test DB")
-	}
-	defer coredb.ResetTestDB(t, testDB)
+	t.Parallel()
+
+	testDB := database.NewTestDatabase(t)
+	testPublishDB := New(testDB)
 	ctx := context.Background()
 
 	// Insert some Exposures.
@@ -129,7 +112,7 @@ func TestExposures(t *testing.T) {
 			nil,
 		},
 	} {
-		got, err := listExposures(ctx, test.criteria)
+		got, err := listExposures(ctx, testPublishDB, test.criteria)
 		if err != nil {
 			t.Fatalf("%+v: %v", test.criteria, err)
 		}
@@ -151,7 +134,7 @@ func TestExposures(t *testing.T) {
 	if gotN != wantN {
 		t.Errorf("DeleteExposures: deleted %d, want %d", gotN, wantN)
 	}
-	got, err := listExposures(ctx, IterateExposuresCriteria{})
+	got, err := listExposures(ctx, testPublishDB, IterateExposuresCriteria{})
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -161,24 +144,24 @@ func TestExposures(t *testing.T) {
 
 }
 
-func listExposures(ctx context.Context, c IterateExposuresCriteria) (_ []*model.Exposure, err error) {
+func listExposures(ctx context.Context, db *PublishDB, c IterateExposuresCriteria) (_ []*model.Exposure, err error) {
 	var exps []*model.Exposure
-	_, err = testPublishDB.IterateExposures(ctx, c, func(e *model.Exposure) error {
+	if _, err := db.IterateExposures(ctx, c, func(e *model.Exposure) error {
 		exps = append(exps, e)
 		return nil
-	})
-	if err != nil {
-		return nil, err
+	}); err != nil {
+		return nil, fmt.Errorf("failed to list exposures: %w", err)
 	}
 	return exps, nil
 }
 
 func TestIterateExposuresCursor(t *testing.T) {
-	if testDB == nil {
-		t.Skip("no test DB")
-	}
-	defer coredb.ResetTestDB(t, testDB)
+	t.Parallel()
+
+	testDB := database.NewTestDatabase(t)
+	testPublishDB := New(testDB)
 	ctx, cancel := context.WithCancel(context.Background())
+
 	// Insert some Exposures.
 	exposures := []*model.Exposure{
 		{