feat: support adding wait strategies as functional option (testcontainers#3161)

mdelapenya · web-flow · commit 2740634ac442 · 2025-05-15T22:50:51.000+02:00
* feat: add a new functional option to append wait strategies to modules

* chore: update modules

* docs: make it mkore clear
diff --git a/docs/features/common_functional_options.md b/docs/features/common_functional_options.md
@@ -242,12 +242,14 @@ You can also use `testcontainers.WithAdditionalLifecycleHooks`, which appends th
 
 #### Wait Strategies
 
-If you need to set a different wait strategy for the container, you can use `testcontainers.WithWaitStrategy` with a valid wait strategy.
+If you need to set a different wait strategy for the container, replacing the existing one, you can use `testcontainers.WithWaitStrategy` with a valid wait strategy.
 
 !!!info
     The default deadline for the wait strategy is 60 seconds.
 
-At the same time, it's possible to set a wait strategy and a custom deadline with `testcontainers.WithWaitStrategyAndDeadline`.
+At the same time, it's possible to replace the wait strategy with a new one and a custom deadline, using `testcontainers.WithWaitStrategyAndDeadline`.
+
+Finally, you can also append a wait strategy to the existing wait strategy, using `testcontainers.WithAdditionalWaitStrategy` and `testcontainers.WithAdditionalWaitStrategyAndDeadline`.
 
 #### Startup Commands
 
diff --git a/docs/modules/index.md b/docs/modules/index.md
@@ -210,8 +210,10 @@ In order to simplify the creation of the container for a given module, `Testcont
 - `testcontainers.WithAdditionalLifecycleHooks`: a function that appends lifecycle hooks to the existing ones for the container request.
 - `testcontainers.WithAlwaysPull`: a function that pulls the image before starting the container.
 - `testcontainers.WithImagePlatform`: a function that sets the image platform for the container request.
-- `testcontainers.WithWaitStrategy`: a function that sets the wait strategy for the container request.
-- `testcontainers.WithWaitStrategyAndDeadline`: a function that sets the wait strategy for the container request with a deadline.
+- `testcontainers.WithWaitStrategy`: a function that replaces the wait strategy for the container request.
+- `testcontainers.WithAdditionalWaitStrategy`: a function that appends the wait strategy for the container request.
+- `testcontainers.WithWaitStrategyAndDeadline`: a function that replaces the wait strategy for the container request with a deadline.
+- `testcontainers.WithAdditionalWaitStrategyAndDeadline`: a function that appends the wait strategy for the container request with a deadline.
 - `testcontainers.WithStartupCommand`: a function that sets the execution of a command when the container starts.
 - `testcontainers.WithAfterReadyCommand`: a function that sets the execution of a command right after the container is ready (its wait strategy is satisfied).
 - `testcontainers.WithDockerfile`: a function that sets the build from a Dockerfile for the container request.
diff --git a/modules/k3s/k3s_test.go b/modules/k3s/k3s_test.go
@@ -146,7 +146,7 @@ func Test_WithManifestOption(t *testing.T) {
 	k3sContainer, err := k3s.Run(ctx,
 		"rancher/k3s:v1.27.1-k3s1",
 		k3s.WithManifest("nginx-manifest.yaml"),
-		testcontainers.WithWaitStrategy(wait.ForExec([]string{"kubectl", "wait", "pod", "nginx", "--for=condition=Ready"})),
+		testcontainers.WithAdditionalWaitStrategy(wait.ForExec([]string{"kubectl", "wait", "pod", "nginx", "--for=condition=Ready"})),
 	)
 	testcontainers.CleanupContainer(t, k3sContainer)
 	require.NoError(t, err)
diff --git a/modules/localstack/examples_test.go b/modules/localstack/examples_test.go
@@ -101,7 +101,7 @@ func ExampleRun_legacyMode() {
 		ctx,
 		"localstack/localstack:0.10.0",
 		testcontainers.WithEnv(map[string]string{"SERVICES": "s3,sqs"}),
-		testcontainers.WithWaitStrategy(wait.ForLog("Ready.").WithStartupTimeout(5*time.Minute).WithOccurrence(1)),
+		testcontainers.WithAdditionalWaitStrategy(wait.ForLog("Ready.").WithStartupTimeout(5*time.Minute).WithOccurrence(1)),
 	)
 	defer func() {
 		if err := testcontainers.TerminateContainer(ctr); err != nil {
diff --git a/modules/mongodb/examples_test.go b/modules/mongodb/examples_test.go
@@ -11,7 +11,6 @@ import (
 
 	"github.com/testcontainers/testcontainers-go"
 	"github.com/testcontainers/testcontainers-go/modules/mongodb"
-	"github.com/testcontainers/testcontainers-go/wait"
 )
 
 func ExampleRun() {
@@ -89,7 +88,6 @@ func ExampleRun_withCredentials() {
 		"mongo:6",
 		mongodb.WithUsername("root"),
 		mongodb.WithPassword("password"),
-		testcontainers.WithWaitStrategy(wait.ForLog("Waiting for connections")),
 	)
 	defer func() {
 		if err := testcontainers.TerminateContainer(ctr); err != nil {
diff --git a/modules/mssql/mssql_test.go b/modules/mssql/mssql_test.go
@@ -50,7 +50,7 @@ func TestMSSQLServerWithMissingEulaOption(t *testing.T) {
 	t.Run("empty", func(t *testing.T) {
 		ctr, err := mssql.Run(ctx,
 			"mcr.microsoft.com/mssql/server:2022-CU14-ubuntu-22.04",
-			testcontainers.WithWaitStrategy(
+			testcontainers.WithAdditionalWaitStrategy(
 				wait.ForLog("The SQL Server End-User License Agreement (EULA) must be accepted")),
 		)
 		testcontainers.CleanupContainer(t, ctr)
@@ -61,7 +61,7 @@ func TestMSSQLServerWithMissingEulaOption(t *testing.T) {
 		ctr, err := mssql.Run(ctx,
 			"mcr.microsoft.com/mssql/server:2022-CU14-ubuntu-22.04",
 			testcontainers.WithEnv(map[string]string{"ACCEPT_EULA": "yes"}),
-			testcontainers.WithWaitStrategy(
+			testcontainers.WithAdditionalWaitStrategy(
 				wait.ForLog("The SQL Server End-User License Agreement (EULA) must be accepted")),
 		)
 		testcontainers.CleanupContainer(t, ctr)
diff --git a/modules/postgres/examples_test.go b/modules/postgres/examples_test.go
@@ -5,11 +5,9 @@ import (
 	"fmt"
 	"log"
 	"path/filepath"
-	"time"
 
 	"github.com/testcontainers/testcontainers-go"
 	"github.com/testcontainers/testcontainers-go/modules/postgres"
-	"github.com/testcontainers/testcontainers-go/wait"
 )
 
 func ExampleRun() {
@@ -27,10 +25,7 @@ func ExampleRun() {
 		postgres.WithDatabase(dbName),
 		postgres.WithUsername(dbUser),
 		postgres.WithPassword(dbPassword),
-		testcontainers.WithWaitStrategy(
-			wait.ForLog("database system is ready to accept connections").
-				WithOccurrence(2).
-				WithStartupTimeout(5*time.Second)),
+		postgres.BasicWaitStrategies(),
 	)
 	defer func() {
 		if err := testcontainers.TerminateContainer(postgresContainer); err != nil {
diff --git a/modules/postgres/postgres_test.go b/modules/postgres/postgres_test.go
@@ -155,7 +155,7 @@ func TestContainerWithWaitForSQL(t *testing.T) {
 			postgres.WithDatabase(dbname),
 			postgres.WithUsername(user),
 			postgres.WithPassword(password),
-			testcontainers.WithWaitStrategy(wait.ForSQL(nat.Port(port), "postgres", dbURL)),
+			testcontainers.WithAdditionalWaitStrategy(wait.ForSQL(nat.Port(port), "postgres", dbURL)),
 		)
 		testcontainers.CleanupContainer(t, ctr)
 		require.NoError(t, err)
@@ -168,7 +168,7 @@ func TestContainerWithWaitForSQL(t *testing.T) {
 			postgres.WithDatabase(dbname),
 			postgres.WithUsername(user),
 			postgres.WithPassword(password),
-			testcontainers.WithWaitStrategy(wait.ForSQL(nat.Port(port), "postgres", dbURL).WithStartupTimeout(time.Second*5).WithQuery("SELECT 10")),
+			testcontainers.WithAdditionalWaitStrategy(wait.ForSQL(nat.Port(port), "postgres", dbURL).WithStartupTimeout(time.Second*5).WithQuery("SELECT 10")),
 		)
 		testcontainers.CleanupContainer(t, ctr)
 		require.NoError(t, err)
@@ -181,7 +181,7 @@ func TestContainerWithWaitForSQL(t *testing.T) {
 			postgres.WithDatabase(dbname),
 			postgres.WithUsername(user),
 			postgres.WithPassword(password),
-			testcontainers.WithWaitStrategy(wait.ForSQL(nat.Port(port), "postgres", dbURL).WithStartupTimeout(time.Second*5).WithQuery("SELECT 'a' from b")),
+			testcontainers.WithAdditionalWaitStrategy(wait.ForSQL(nat.Port(port), "postgres", dbURL).WithStartupTimeout(time.Second*5).WithQuery("SELECT 'a' from b")),
 		)
 		testcontainers.CleanupContainer(t, ctr)
 		require.Error(t, err)
@@ -225,7 +225,7 @@ func TestWithSSL(t *testing.T) {
 		postgres.WithDatabase(dbname),
 		postgres.WithUsername(user),
 		postgres.WithPassword(password),
-		testcontainers.WithWaitStrategy(wait.ForLog("database system is ready to accept connections").WithOccurrence(2).WithStartupTimeout(5*time.Second)),
+		testcontainers.WithAdditionalWaitStrategy(wait.ForLog("database system is ready to accept connections").WithOccurrence(2).WithStartupTimeout(5*time.Second)),
 		postgres.WithSSLCert(caCert.CertPath, serverCerts.CertPath, serverCerts.KeyPath),
 	)
 
@@ -255,7 +255,7 @@ func TestSSLValidatesKeyMaterialPath(t *testing.T) {
 		postgres.WithDatabase(dbname),
 		postgres.WithUsername(user),
 		postgres.WithPassword(password),
-		testcontainers.WithWaitStrategy(wait.ForLog("database system is ready to accept connections").WithOccurrence(2).WithStartupTimeout(5*time.Second)),
+		testcontainers.WithAdditionalWaitStrategy(wait.ForLog("database system is ready to accept connections").WithOccurrence(2).WithStartupTimeout(5*time.Second)),
 		postgres.WithSSLCert("", "", ""),
 	)
 
diff --git a/modules/postgres/wait_strategies.go b/modules/postgres/wait_strategies.go
@@ -14,7 +14,7 @@ import (
 //     Without this, the tests will be flaky on those OSes!
 func BasicWaitStrategies() testcontainers.CustomizeRequestOption {
 	// waitStrategy {
-	return testcontainers.WithWaitStrategy(
+	return testcontainers.WithAdditionalWaitStrategy(
 		// First, we wait for the container to log readiness twice.
 		// This is because it will restart itself after the first startup.
 		wait.ForLog("database system is ready to accept connections").WithOccurrence(2),
diff --git a/options.go b/options.go
@@ -373,12 +373,17 @@ func WithAfterReadyCommand(execs ...Executable) CustomizeRequestOption {
 	}
 }
 
-// WithWaitStrategy sets the wait strategy for a container, using 60 seconds as deadline
+// WithWaitStrategy replaces the wait strategy for a container, using 60 seconds as deadline
 func WithWaitStrategy(strategies ...wait.Strategy) CustomizeRequestOption {
 	return WithWaitStrategyAndDeadline(60*time.Second, strategies...)
 }
 
-// WithWaitStrategyAndDeadline sets the wait strategy for a container, including deadline
+// WithAdditionalWaitStrategy appends the wait strategy for a container, using 60 seconds as deadline
+func WithAdditionalWaitStrategy(strategies ...wait.Strategy) CustomizeRequestOption {
+	return WithAdditionalWaitStrategyAndDeadline(60*time.Second, strategies...)
+}
+
+// WithWaitStrategyAndDeadline replaces the wait strategy for a container, including deadline
 func WithWaitStrategyAndDeadline(deadline time.Duration, strategies ...wait.Strategy) CustomizeRequestOption {
 	return func(req *GenericContainerRequest) error {
 		req.WaitingFor = wait.ForAll(strategies...).WithDeadline(deadline)
@@ -387,6 +392,24 @@ func WithWaitStrategyAndDeadline(deadline time.Duration, strategies ...wait.Stra
 	}
 }
 
+// WithAdditionalWaitStrategyAndDeadline appends the wait strategy for a container, including deadline
+func WithAdditionalWaitStrategyAndDeadline(deadline time.Duration, strategies ...wait.Strategy) CustomizeRequestOption {
+	return func(req *GenericContainerRequest) error {
+		if req.WaitingFor == nil {
+			req.WaitingFor = wait.ForAll(strategies...).WithDeadline(deadline)
+			return nil
+		}
+
+		wss := make([]wait.Strategy, 0, len(strategies)+1)
+		wss = append(wss, req.WaitingFor)
+		wss = append(wss, strategies...)
+
+		req.WaitingFor = wait.ForAll(wss...).WithDeadline(deadline)
+
+		return nil
+	}
+}
+
 // WithImageMount mounts an image to a container, passing the source image name,
 // the relative subpath to mount in that image, and the mount point in the target container.
 // This option validates that the subpath is a relative path, raising an error otherwise.
diff --git a/options_test.go b/options_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"io"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 
@@ -784,3 +785,124 @@ func TestWithNoStart(t *testing.T) {
 	require.NoError(t, err)
 	require.False(t, req.Started)
 }
+
+func TestWithWaitStrategy(t *testing.T) {
+	testDuration := 10 * time.Second
+	defaultDuration := 60 * time.Second
+
+	waitForFoo := wait.ForLog("foo")
+	waitForBar := wait.ForLog("bar")
+
+	testWaitFor := func(t *testing.T, replace bool, customDuration *time.Duration, initial wait.Strategy, add wait.Strategy, expected wait.Strategy) {
+		t.Helper()
+
+		req := &testcontainers.GenericContainerRequest{
+			ContainerRequest: testcontainers.ContainerRequest{
+				WaitingFor: initial,
+			},
+		}
+
+		var opt testcontainers.CustomizeRequestOption
+		if replace {
+			opt = testcontainers.WithWaitStrategy(add)
+			if customDuration != nil {
+				opt = testcontainers.WithWaitStrategyAndDeadline(*customDuration, add)
+			}
+		} else {
+			opt = testcontainers.WithAdditionalWaitStrategy(add)
+			if customDuration != nil {
+				opt = testcontainers.WithAdditionalWaitStrategyAndDeadline(*customDuration, add)
+			}
+		}
+		require.NoError(t, opt.Customize(req))
+		require.Equal(t, expected, req.WaitingFor)
+	}
+
+	t.Run("replace-nil", func(t *testing.T) {
+		t.Run("default-duration", func(t *testing.T) {
+			testWaitFor(t,
+				true,
+				nil,
+				nil,
+				waitForFoo,
+				wait.ForAll(waitForFoo).WithDeadline(defaultDuration),
+			)
+		})
+
+		t.Run("custom-duration", func(t *testing.T) {
+			testWaitFor(t,
+				true,
+				&testDuration,
+				nil,
+				waitForFoo,
+				wait.ForAll(waitForFoo).WithDeadline(testDuration),
+			)
+		})
+	})
+
+	t.Run("replace-existing", func(t *testing.T) {
+		t.Run("default-duration", func(t *testing.T) {
+			testWaitFor(t,
+				true,
+				nil,
+				waitForFoo,
+				waitForBar,
+				wait.ForAll(waitForBar).WithDeadline(defaultDuration),
+			)
+		})
+
+		t.Run("custom-duration", func(t *testing.T) {
+			testWaitFor(t,
+				true,
+				&testDuration,
+				waitForFoo,
+				waitForBar,
+				wait.ForAll(waitForBar).WithDeadline(testDuration),
+			)
+		})
+	})
+
+	t.Run("add-to-nil", func(t *testing.T) {
+		t.Run("default-duration", func(t *testing.T) {
+			testWaitFor(t,
+				false,
+				nil,
+				nil,
+				waitForFoo,
+				wait.ForAll(waitForFoo).WithDeadline(defaultDuration),
+			)
+		})
+
+		t.Run("custom-duration", func(t *testing.T) {
+			testWaitFor(t,
+				false,
+				&testDuration,
+				nil,
+				waitForFoo,
+				wait.ForAll(waitForFoo).WithDeadline(testDuration),
+			)
+		})
+	})
+
+	t.Run("add-to-existing", func(t *testing.T) {
+		t.Run("default-duration", func(t *testing.T) {
+			testWaitFor(t,
+				false,
+				nil,
+				waitForFoo,
+				waitForBar,
+				wait.ForAll(waitForFoo, waitForBar).WithDeadline(defaultDuration),
+			)
+		})
+
+		t.Run("custom-duration", func(t *testing.T) {
+			testWaitFor(t,
+				false,
+				&testDuration,
+				waitForFoo,
+				waitForBar,
+				wait.ForAll(waitForFoo, waitForBar).WithDeadline(testDuration),
+			)
+		})
+	})
+}

Original file line number	Diff line number	Diff line change
`@@ -146,7 +146,7 @@ func Test_WithManifestOption(t *testing.T) {`
`146`	`146`	`k3sContainer, err := k3s.Run(ctx,`
`147`	`147`	`"rancher/k3s:v1.27.1-k3s1",`
`148`	`148`	`k3s.WithManifest("nginx-manifest.yaml"),`
`149`		`- testcontainers.WithWaitStrategy(wait.ForExec([]string{"kubectl", "wait", "pod", "nginx", "--for=condition=Ready"})),`
	`149`	`+ testcontainers.WithAdditionalWaitStrategy(wait.ForExec([]string{"kubectl", "wait", "pod", "nginx", "--for=condition=Ready"})),`
`150`	`150`	`)`
`151`	`151`	`testcontainers.CleanupContainer(t, k3sContainer)`
`152`	`152`	`require.NoError(t, err)`
Original file line number	Diff line number	Diff line change
`@@ -101,7 +101,7 @@ func ExampleRun_legacyMode() {`
`101`	`101`	`ctx,`
`102`	`102`	`"localstack/localstack:0.10.0",`
`103`	`103`	`testcontainers.WithEnv(map[string]string{"SERVICES": "s3,sqs"}),`
`104`		`- testcontainers.WithWaitStrategy(wait.ForLog("Ready.").WithStartupTimeout(5*time.Minute).WithOccurrence(1)),`
	`104`	`+ testcontainers.WithAdditionalWaitStrategy(wait.ForLog("Ready.").WithStartupTimeout(5*time.Minute).WithOccurrence(1)),`
`105`	`105`	`)`
`106`	`106`	`defer func() {`
`107`	`107`	`if err := testcontainers.TerminateContainer(ctr); err != nil {`
Original file line number	Diff line number	Diff line change
`@@ -11,7 +11,6 @@ import (`
`11`	`11`
`12`	`12`	`"github.com/testcontainers/testcontainers-go"`
`13`	`13`	`"github.com/testcontainers/testcontainers-go/modules/mongodb"`
`14`		`- "github.com/testcontainers/testcontainers-go/wait"`
`15`	`14`	`)`
`16`	`15`
`17`	`16`	`func ExampleRun() {`
`@@ -89,7 +88,6 @@ func ExampleRun_withCredentials() {`
`89`	`88`	`"mongo:6",`
`90`	`89`	`mongodb.WithUsername("root"),`
`91`	`90`	`mongodb.WithPassword("password"),`
`92`		`- testcontainers.WithWaitStrategy(wait.ForLog("Waiting for connections")),`
`93`	`91`	`)`
`94`	`92`	`defer func() {`
`95`	`93`	`if err := testcontainers.TerminateContainer(ctr); err != nil {`