Clean up container on connection failure, switch to ory/dockertest on package postgresql (#5050)

api/api_integration_test.go

@@ -145,6 +145,12 @@ func testPostgresDB(t testing.TB) (string, func()) {
 		t.Fatalf("postgresdb: could not start container: %s", err)
 	}
 
+	cleanup := func() {
+		if err := pool.Purge(resource); err != nil {
+			t.Fatalf("failed to cleanup local container: %s", err)
+		}
+	}
+
 	addr := fmt.Sprintf("postgres://postgres:secret@localhost:%s/database?sslmode=disable", resource.GetPort("5432/tcp"))
 
 	if err := pool.Retry(func() error {
@@ -155,12 +161,9 @@ func testPostgresDB(t testing.TB) (string, func()) {
 		defer db.Close()
 		return db.Ping()
 	}); err != nil {
+		cleanup()
 		t.Fatalf("postgresdb: could not connect: %s", err)
 	}
 
-	return addr, func() {
-		if err := pool.Purge(resource); err != nil {
-			t.Fatalf("postgresdb: failed to cleanup container: %s", err)
-		}
-	}
+	return addr, cleanup
 }

builtin/logical/database/backend_test.go

@@ -77,6 +77,7 @@ func preparePostgresTestContainer(t *testing.T, s logical.Storage, b logical.Bac
 
 		return nil
 	}); err != nil {
+		cleanup()
 		t.Fatalf("Could not connect to PostgreSQL docker container: %s", err)
 	}
 

builtin/logical/postgresql/backend_test.go

@@ -9,71 +9,57 @@ import (
 	"os"
 	"path"
 	"reflect"
-	"sync"
 	"testing"
-	"time"
 
 	"github.com/hashicorp/vault/logical"
 	logicaltest "github.com/hashicorp/vault/logical/testing"
 	"github.com/lib/pq"
 	"github.com/mitchellh/mapstructure"
-	dockertest "gopkg.in/ory-am/dockertest.v2"
+	"github.com/ory/dockertest"
 )
 
-var (
-	testImagePull sync.Once
-)
-
-func prepareTestContainer(t *testing.T, s logical.Storage, b logical.Backend) (cid dockertest.ContainerID, retURL string) {
+func prepareTestContainer(t *testing.T) (cleanup func(), retURL string) {
 	if os.Getenv("PG_URL") != "" {
-		return "", os.Getenv("PG_URL")
+		return func() {}, os.Getenv("PG_URL")
 	}
 
-	// Without this the checks for whether the container has started seem to
-	// never actually pass. There's really no reason to expose the test
-	// containers, so don't.
-	dockertest.BindDockerToLocalhost = "yep"
-
-	testImagePull.Do(func() {
-		dockertest.Pull("postgres")
-	})
-
-	cid, connErr := dockertest.ConnectToPostgreSQL(60, 500*time.Millisecond, func(connURL string) bool {
-		// This will cause a validation to run
-		resp, err := b.HandleRequest(context.Background(), &logical.Request{
-			Storage:   s,
-			Operation: logical.UpdateOperation,
-			Path:      "config/connection",
-			Data: map[string]interface{}{
-				"connection_url": connURL,
-			},
-		})
-		if err != nil || (resp != nil && resp.IsError()) {
-			// It's likely not up and running yet, so return false and try again
-			return false
-		}
-		if resp == nil {
-			t.Fatal("expected warning")
+	pool, err := dockertest.NewPool("")
+	if err != nil {
+		t.Fatalf("Failed to connect to docker: %s", err)
 	}
 
-		retURL = connURL
-		return true
-	})
+	resource, err := pool.Run("postgres", "latest", []string{"POSTGRES_PASSWORD=secret", "POSTGRES_DB=database"})
+	if err != nil {
+		t.Fatalf("Could not start local PostgreSQL docker container: %s", err)
+	}
 
-	if connErr != nil {
-		t.Fatalf("could not connect to database: %v", connErr)
+	cleanup = func() {
+		err := pool.Purge(resource)
+		if err != nil {
+			t.Fatalf("Failed to cleanup local container: %s", err)
+		}
+	}
+
+	retURL = fmt.Sprintf("postgres://postgres:secret@localhost:%s/database?sslmode=disable", resource.GetPort("5432/tcp"))
+
+	// exponential backoff-retry
+	if err = pool.Retry(func() error {
+		var err error
+		var db *sql.DB
+		db, err = sql.Open("postgres", retURL)
+		if err != nil {
+			return err
+		}
+		defer db.Close()
+		return db.Ping()
+	}); err != nil {
+		cleanup()
+		t.Fatalf("Could not connect to PostgreSQL docker container: %s", err)
 	}
 
 	return
 }
 
-func cleanupTestContainer(t *testing.T, cid dockertest.ContainerID) {
-	err := cid.KillRemove()
-	if err != nil {
-		t.Fatal(err)
-	}
-}
-
 func TestBackend_config_connection(t *testing.T) {
 	var resp *logical.Response
 	var err error
@@ -123,14 +109,12 @@ func TestBackend_basic(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	cid, connURL := prepareTestContainer(t, config.StorageView, b)
-	if cid != "" {
-		defer cleanupTestContainer(t, cid)
-	}
+	cleanup, connURL := prepareTestContainer(t)
+	defer cleanup()
+
 	connData := map[string]interface{}{
 		"connection_url": connURL,
 	}
-
 	logicaltest.Test(t, logicaltest.TestCase{
 		Backend: b,
 		Steps: []logicaltest.TestStep{
@@ -149,14 +133,12 @@ func TestBackend_roleCrud(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	cid, connURL := prepareTestContainer(t, config.StorageView, b)
-	if cid != "" {
-		defer cleanupTestContainer(t, cid)
-	}
+	cleanup, connURL := prepareTestContainer(t)
+	defer cleanup()
+
 	connData := map[string]interface{}{
 		"connection_url": connURL,
 	}
-
 	logicaltest.Test(t, logicaltest.TestCase{
 		Backend: b,
 		Steps: []logicaltest.TestStep{
@@ -177,14 +159,12 @@ func TestBackend_BlockStatements(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	cid, connURL := prepareTestContainer(t, config.StorageView, b)
-	if cid != "" {
-		defer cleanupTestContainer(t, cid)
-	}
+	cleanup, connURL := prepareTestContainer(t)
+	defer cleanup()
+
 	connData := map[string]interface{}{
 		"connection_url": connURL,
 	}
-
 	jsonBlockStatement, err := json.Marshal(testBlockStatementRoleSlice)
 	if err != nil {
 		t.Fatal(err)
@@ -209,14 +189,12 @@ func TestBackend_roleReadOnly(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	cid, connURL := prepareTestContainer(t, config.StorageView, b)
-	if cid != "" {
-		defer cleanupTestContainer(t, cid)
-	}
+	cleanup, connURL := prepareTestContainer(t)
+	defer cleanup()
+
 	connData := map[string]interface{}{
 		"connection_url": connURL,
 	}
-
 	logicaltest.Test(t, logicaltest.TestCase{
 		Backend: b,
 		Steps: []logicaltest.TestStep{
@@ -242,14 +220,12 @@ func TestBackend_roleReadOnly_revocationSQL(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	cid, connURL := prepareTestContainer(t, config.StorageView, b)
-	if cid != "" {
-		defer cleanupTestContainer(t, cid)
-	}
+	cleanup, connURL := prepareTestContainer(t)
+	defer cleanup()
+
 	connData := map[string]interface{}{
 		"connection_url": connURL,
 	}
-
 	logicaltest.Test(t, logicaltest.TestCase{
 		Backend: b,
 		Steps: []logicaltest.TestStep{

physical/dynamodb/dynamodb_test.go

@@ -311,6 +311,7 @@ func prepareDynamoDBTestContainer(t *testing.T) (cleanup func(), retAddress stri
 		}
 		return nil
 	}); err != nil {
+		cleanup()
 		t.Fatalf("Could not connect to docker: %s", err)
 	}
 	return cleanup, retAddress, credentials.NewStaticCredentials("fake", "fake", "")

plugins/database/mongodb/mongodb_test.go

@@ -60,6 +60,7 @@ func prepareMongoDBTestContainer(t *testing.T) (cleanup func(), retURL string) {
 		session.SetSocketTimeout(1 * time.Minute)
 		return session.Ping()
 	}); err != nil {
+		cleanup()
 		t.Fatalf("Could not connect to mongo docker container: %s", err)
 	}
 

plugins/database/postgresql/postgresql_test.go

@@ -53,6 +53,7 @@ func preparePostgresTestContainer(t *testing.T) (cleanup func(), retURL string)
 		defer db.Close()
 		return db.Ping()
 	}); err != nil {
+		cleanup()
 		t.Fatalf("Could not connect to PostgreSQL docker container: %s", err)
 	}