fix(cockroachdb): Fixes cockroachdb wait strategy handling (#2456)

* Fixes cockroachdb wait strategy handling * Clarify test intent * Merge in default strategy to get connection * Use subtests * Separate out Connection logic * Feedback from PR * Honor default settings * fix: make lint --------- Co-authored-by: bstrausser <bstrausser@locusrobotics.com> Co-authored-by: Manuel de la Peña <mdelapenya@gmail.com>
testcontainers · Jun 11, 2024 · d483772 · d483772
1 parent 0662f1f
commit d483772
Show file tree

Hide file tree

Showing 2 changed files with 106 additions and 16 deletions.
diff --git a/modules/cockroachdb/cockroachdb.go b/modules/cockroachdb/cockroachdb.go
@@ -169,25 +169,32 @@ func addWaitingFor(req *testcontainers.GenericContainerRequest, opts options) er
 		tlsConfig = cfg
 	}
 
-	req.WaitingFor = wait.ForAll(
+	sqlWait := wait.ForSQL(defaultSQLPort, "pgx/v5", func(host string, port nat.Port) string {
+		connStr := connString(opts, host, port)
+		if tlsConfig == nil {
+			return connStr
+		}
+
+		// register TLS config with pgx driver
+		connCfg, err := pgx.ParseConfig(connStr)
+		if err != nil {
+			panic(err)
+		}
+		connCfg.TLSConfig = tlsConfig
+
+		return stdlib.RegisterConnConfig(connCfg)
+	})
+	defaultStrategy := wait.ForAll(
 		wait.ForHTTP("/health").WithPort(defaultAdminPort),
-		wait.ForSQL(defaultSQLPort, "pgx/v5", func(host string, port nat.Port) string {
-			connStr := connString(opts, host, port)
-			if tlsConfig == nil {
-				return connStr
-			}
-
-			// register TLS config with pgx driver
-			connCfg, err := pgx.ParseConfig(connStr)
-			if err != nil {
-				panic(err)
-			}
-			connCfg.TLSConfig = tlsConfig
-
-			return stdlib.RegisterConnConfig(connCfg)
-		}),
+		sqlWait,
 	)
 
+	if req.WaitingFor == nil {
+		req.WaitingFor = defaultStrategy
+	} else {
+		req.WaitingFor = wait.ForAll(req.WaitingFor, defaultStrategy)
+	}
+
 	return nil
 }
 

diff --git a/modules/cockroachdb/cockroachdb_test.go b/modules/cockroachdb/cockroachdb_test.go
@@ -6,13 +6,15 @@ import (
 	"net/url"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/jackc/pgx/v5"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 
 	"github.com/testcontainers/testcontainers-go"
 	"github.com/testcontainers/testcontainers-go/modules/cockroachdb"
+	"github.com/testcontainers/testcontainers-go/wait"
 )
 
 func TestCockroach_Insecure(t *testing.T) {
@@ -144,6 +146,87 @@ func (suite *AuthNSuite) TestQuery() {
 	suite.Equal(523123, id)
 }
 
+// TestWithWaitStrategyAndDeadline covers a previous regression, container creation needs to fail to cover that path.
+func (suite *AuthNSuite) TestWithWaitStrategyAndDeadline() {
+	nodeStartUpCompleted := "node startup completed"
+
+	suite.Run("Expected Failure To Run", func() {
+		ctx := context.Background()
+
+		// This will never match a log statement
+		suite.opts = append(suite.opts, testcontainers.WithWaitStrategyAndDeadline(time.Millisecond*250, wait.ForLog("Won't Exist In Logs")))
+		container, err := cockroachdb.RunContainer(ctx, suite.opts...)
+
+		suite.Require().ErrorIs(err, context.DeadlineExceeded)
+		suite.T().Cleanup(func() {
+			if container != nil {
+				err := container.Terminate(ctx)
+				suite.Require().NoError(err)
+			}
+		})
+	})
+
+	suite.Run("Expected Failure To Run But Would Succeed ", func() {
+		ctx := context.Background()
+
+		// This will timeout as we didn't give enough time for intialization, but would have succeeded otherwise
+		suite.opts = append(suite.opts, testcontainers.WithWaitStrategyAndDeadline(time.Millisecond*20, wait.ForLog(nodeStartUpCompleted)))
+		container, err := cockroachdb.RunContainer(ctx, suite.opts...)
+
+		suite.Require().ErrorIs(err, context.DeadlineExceeded)
+		suite.T().Cleanup(func() {
+			if container != nil {
+				err := container.Terminate(ctx)
+				suite.Require().NoError(err)
+			}
+		})
+	})
+
+	suite.Run("Succeeds And Executes Commands", func() {
+		ctx := context.Background()
+
+		// This will succeed
+		suite.opts = append(suite.opts, testcontainers.WithWaitStrategyAndDeadline(time.Second*60, wait.ForLog(nodeStartUpCompleted)))
+		container, err := cockroachdb.RunContainer(ctx, suite.opts...)
+		suite.Require().NoError(err)
+
+		conn, err := conn(ctx, container)
+		suite.Require().NoError(err)
+		defer conn.Close(ctx)
+
+		_, err = conn.Exec(ctx, "CREATE TABLE test (id INT PRIMARY KEY)")
+		suite.Require().NoError(err)
+		suite.T().Cleanup(func() {
+			if container != nil {
+				err := container.Terminate(ctx)
+				suite.Require().NoError(err)
+			}
+		})
+	})
+
+	suite.Run("Succeeds And Executes Commands Waiting on HTTP Endpoint", func() {
+		ctx := context.Background()
+
+		// This will succeed
+		suite.opts = append(suite.opts, testcontainers.WithWaitStrategyAndDeadline(time.Second*60, wait.ForHTTP("/health").WithPort("8080/tcp")))
+		container, err := cockroachdb.RunContainer(ctx, suite.opts...)
+		suite.Require().NoError(err)
+
+		conn, err := conn(ctx, container)
+		suite.Require().NoError(err)
+		defer conn.Close(ctx)
+
+		_, err = conn.Exec(ctx, "CREATE TABLE test (id INT PRIMARY KEY)")
+		suite.Require().NoError(err)
+		suite.T().Cleanup(func() {
+			if container != nil {
+				err := container.Terminate(ctx)
+				suite.Require().NoError(err)
+			}
+		})
+	})
+}
+
 func conn(ctx context.Context, container *cockroachdb.CockroachDBContainer) (*pgx.Conn, error) {
 	cfg, err := pgx.ParseConfig(container.MustConnectionString(ctx))
 	if err != nil {