review: move Strategy#Timeout to StrategyTimeout interface

testcontainers · Nov 20, 2022 · c594243 · c594243
1 parent 1f5f932
commit c594243
Show file tree

Hide file tree

Showing 10 changed files with 31 additions and 8 deletions.
diff --git a/wait/all.go b/wait/all.go
@@ -8,6 +8,7 @@ import (
 
 // Implement interface
 var _ Strategy = (*MultiStrategy)(nil)
+var _ StrategyTimeout = (*MultiStrategy)(nil)
 
 type MultiStrategy struct {
 	// all Strategies should have a startupTimeout to avoid waiting infinitely
@@ -47,7 +48,7 @@ func (ms *MultiStrategy) Timeout() *time.Duration {
 	return ms.timeout
 }
 
-func (ms *MultiStrategy) WaitUntilReady(ctx context.Context, target StrategyTarget) (err error) {
+func (ms *MultiStrategy) WaitUntilReady(ctx context.Context, target StrategyTarget) error {
 	var cancel context.CancelFunc
 	if ms.deadline != nil {
 		ctx, cancel = context.WithTimeout(ctx, *ms.deadline)
@@ -60,15 +61,20 @@ func (ms *MultiStrategy) WaitUntilReady(ctx context.Context, target StrategyTarg
 
 	for _, strategy := range ms.Strategies {
 		strategyCtx := ctx
-		if ms.Timeout() != nil && strategy.Timeout() == nil {
-			strategyCtx, cancel = context.WithTimeout(ctx, *ms.Timeout())
-			defer cancel()
+
+		// Set default Timeout when strategy implements StrategyTimeout
+		if st, ok := strategy.(StrategyTimeout); ok {
+			if ms.Timeout() != nil && st.Timeout() == nil {
+				strategyCtx, cancel = context.WithTimeout(ctx, *ms.Timeout())
+				defer cancel()
+			}
 		}
 
 		err := strategy.WaitUntilReady(strategyCtx, target)
 		if err != nil {
 			return err
 		}
 	}
+
 	return nil
 }
diff --git a/wait/exec.go b/wait/exec.go
@@ -7,6 +7,7 @@ import (
 
 // Implement interface
 var _ Strategy = (*ExecStrategy)(nil)
+var _ StrategyTimeout = (*ExecStrategy)(nil)
 
 type ExecStrategy struct {
 	// all Strategies should have a startupTimeout to avoid waiting infinitely

diff --git a/wait/exit.go b/wait/exit.go
@@ -8,6 +8,7 @@ import (
 
 // Implement interface
 var _ Strategy = (*ExitStrategy)(nil)
+var _ StrategyTimeout = (*ExitStrategy)(nil)
 
 // ExitStrategy will wait until container exit
 type ExitStrategy struct {

diff --git a/wait/health.go b/wait/health.go
@@ -7,6 +7,7 @@ import (
 
 // Implement interface
 var _ Strategy = (*HealthStrategy)(nil)
+var _ StrategyTimeout = (*HealthStrategy)(nil)
 
 // HealthStrategy will wait until the container becomes healthy
 type HealthStrategy struct {

diff --git a/wait/host_port.go b/wait/host_port.go
@@ -14,6 +14,7 @@ import (
 
 // Implement interface
 var _ Strategy = (*HostPortStrategy)(nil)
+var _ StrategyTimeout = (*HostPortStrategy)(nil)
 
 type HostPortStrategy struct {
 	// Port is a string containing port number and protocol in the format "80/tcp"

diff --git a/wait/http.go b/wait/http.go
@@ -17,6 +17,7 @@ import (
 
 // Implement interface
 var _ Strategy = (*HTTPStrategy)(nil)
+var _ StrategyTimeout = (*HTTPStrategy)(nil)
 
 type HTTPStrategy struct {
 	// all Strategies should have a startupTimeout to avoid waiting infinitely

diff --git a/wait/log.go b/wait/log.go
@@ -9,6 +9,7 @@ import (
 
 // Implement interface
 var _ Strategy = (*LogStrategy)(nil)
+var _ StrategyTimeout = (*LogStrategy)(nil)
 
 // LogStrategy will wait until a given log entry shows up in the docker logs
 type LogStrategy struct {

diff --git a/wait/nop.go b/wait/nop.go
@@ -7,9 +7,12 @@ import (
 
 	"github.com/docker/docker/api/types"
 	"github.com/docker/go-connections/nat"
-	tcexec "github.com/testcontainers/testcontainers-go/exec"
+	"github.com/testcontainers/testcontainers-go/exec"
 )
 
+var _ Strategy = (*NopStrategy)(nil)
+var _ StrategyTimeout = (*NopStrategy)(nil)
+
 type NopStrategy struct {
 	timeout        *time.Duration
 	waitUntilReady func(context.Context, StrategyTarget) error
@@ -57,7 +60,7 @@ func (st NopStrategyTarget) Logs(_ context.Context) (io.ReadCloser, error) {
 	return st.ReaderCloser, nil
 }
 
-func (st NopStrategyTarget) Exec(_ context.Context, cmd []string, options ...tcexec.ProcessOption) (int, io.Reader, error) {
+func (st NopStrategyTarget) Exec(_ context.Context, _ []string, _ ...exec.ProcessOption) (int, io.Reader, error) {
 	return 0, nil, nil
 }
 

diff --git a/wait/sql.go b/wait/sql.go
@@ -9,6 +9,9 @@ import (
 	"github.com/docker/go-connections/nat"
 )
 
+var _ Strategy = (*waitForSql)(nil)
+var _ StrategyTimeout = (*waitForSql)(nil)
+
 const defaultForSqlQuery = "SELECT 1"
 
 // ForSQL constructs a new waitForSql strategy for the given driver

diff --git a/wait/wait.go b/wait/wait.go
@@ -7,11 +7,16 @@ import (
 
 	"github.com/docker/docker/api/types"
 	"github.com/docker/go-connections/nat"
-	tcexec "github.com/testcontainers/testcontainers-go/exec"
+	"github.com/testcontainers/testcontainers-go/exec"
 )
 
+// Strategy defines the basic interface for a Wait Strategy
 type Strategy interface {
 	WaitUntilReady(context.Context, StrategyTarget) error
+}
+
+// StrategyTimeout allows MultiStrategy to configure a Strategy's Timeout
+type StrategyTimeout interface {
 	Timeout() *time.Duration
 }
 
@@ -20,7 +25,7 @@ type StrategyTarget interface {
 	Ports(ctx context.Context) (nat.PortMap, error)
 	MappedPort(context.Context, nat.Port) (nat.Port, error)
 	Logs(context.Context) (io.ReadCloser, error)
-	Exec(ctx context.Context, cmd []string, options ...tcexec.ProcessOption) (int, io.Reader, error)
+	Exec(context.Context, []string, ...exec.ProcessOption) (int, io.Reader, error)
 	State(context.Context) (*types.ContainerState, error)
 }