From ece7569aca7eef5c009025a00ef40a5fdb32b07c Mon Sep 17 00:00:00 2001 From: Calvin Leung Huang Date: Tue, 7 Aug 2018 09:56:33 -0700 Subject: [PATCH] Clean up container on connection failure, switch to ory/dockertest on package postgresql (#5050) --- api/api_integration_test.go | 13 +- builtin/logical/database/backend_test.go | 1 + builtin/logical/postgresql/backend_test.go | 116 +++++++----------- physical/dynamodb/dynamodb_test.go | 1 + plugins/database/mongodb/mongodb_test.go | 1 + .../database/postgresql/postgresql_test.go | 1 + 6 files changed, 58 insertions(+), 75 deletions(-) diff --git a/api/api_integration_test.go b/api/api_integration_test.go index bde6155a6c..7d4eb0705f 100644 --- a/api/api_integration_test.go +++ b/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 } diff --git a/builtin/logical/database/backend_test.go b/builtin/logical/database/backend_test.go index 979a2d72d4..69dcb0f554 100644 --- a/builtin/logical/database/backend_test.go +++ b/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) } diff --git a/builtin/logical/postgresql/backend_test.go b/builtin/logical/postgresql/backend_test.go index 8d135028e5..af20ee6797 100644 --- a/builtin/logical/postgresql/backend_test.go +++ b/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" + pool, err := dockertest.NewPool("") + if err != nil { + t.Fatalf("Failed to connect to docker: %s", err) + } - testImagePull.Do(func() { - dockertest.Pull("postgres") - }) + 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) + } - 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 + cleanup = func() { + err := pool.Purge(resource) + if err != nil { + t.Fatalf("Failed to cleanup local container: %s", err) } - if resp == nil { - t.Fatal("expected warning") + } + + 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 } - - retURL = connURL - return true - }) - - if connErr != nil { - t.Fatalf("could not connect to database: %v", connErr) + 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{ diff --git a/physical/dynamodb/dynamodb_test.go b/physical/dynamodb/dynamodb_test.go index 0fadbdcc26..85571586e7 100644 --- a/physical/dynamodb/dynamodb_test.go +++ b/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", "") diff --git a/plugins/database/mongodb/mongodb_test.go b/plugins/database/mongodb/mongodb_test.go index 2a9df469e2..779f312aaa 100644 --- a/plugins/database/mongodb/mongodb_test.go +++ b/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) } diff --git a/plugins/database/postgresql/postgresql_test.go b/plugins/database/postgresql/postgresql_test.go index 24a127bc86..4b4583c381 100644 --- a/plugins/database/postgresql/postgresql_test.go +++ b/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) }