Add wait.ExecStrategy to wait on cmd exec in a container

wait/exec.go

@@ -0,0 +1,77 @@
+package wait
+
+import (
+	"context"
+	"time"
+)
+
+// Implement interface
+var _ Strategy = (*ExecStrategy)(nil)
+
+type ExecStrategy struct {
+	// all Strategies should have a startupTimeout to avoid waiting infinitely
+	startupTimeout time.Duration
+	cmd            []string
+
+	// additional properties
+	ExitCodeMatcher func(exitCode int) bool
+	PollInterval    time.Duration
+}
+
+// NewExecStrategy constructs an Exec strategy ...
+func NewExecStrategy(cmd []string) *ExecStrategy {
+	return &ExecStrategy{
+		startupTimeout:  defaultStartupTimeout(),
+		cmd:             cmd,
+		ExitCodeMatcher: defaultExitCodeMatcher,
+		PollInterval:    defaultPollInterval(),
+	}
+}
+
+func defaultExitCodeMatcher(exitCode int) bool {
+	return exitCode == 0
+}
+
+func (ws *ExecStrategy) WithStartupTimeout(startupTimeout time.Duration) *ExecStrategy {
+	ws.startupTimeout = startupTimeout
+	return ws
+}
+
+func (ws *ExecStrategy) WithExitCodeMatcher(exitCodeMatcher func(exitCode int) bool) *ExecStrategy {
+	ws.ExitCodeMatcher = exitCodeMatcher
+	return ws
+}
+
+// WithPollInterval can be used to override the default polling interval of 100 milliseconds
+func (ws *ExecStrategy) WithPollInterval(pollInterval time.Duration) *ExecStrategy {
+	ws.PollInterval = pollInterval
+	return ws
+}
+
+// ForExec is a convenience method to assign ExecStrategy
+func ForExec(cmd []string) *ExecStrategy {
+	return NewExecStrategy(cmd)
+}
+
+func (ws ExecStrategy) WaitUntilReady(ctx context.Context, target StrategyTarget) error {
+	// limit context to startupTimeout
+	ctx, cancelContext := context.WithTimeout(ctx, ws.startupTimeout)
+	defer cancelContext()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-time.After(ws.PollInterval):
+			exitCode, err := target.Exec(ctx, ws.cmd)
+			if err != nil {
+				return err
+			}
+			if !ws.ExitCodeMatcher(exitCode) {
+				continue
+			}
+
+			return nil
+		}
+	}
+}

wait/exec_test.go

@@ -0,0 +1,130 @@
+package wait_test
+
+import (
+	"context"
+	"errors"
+	"io"
+	"testing"
+	"time"
+
+	"github.com/docker/docker/api/types"
+	"github.com/docker/go-connections/nat"
+
+	"github.com/testcontainers/testcontainers-go"
+	"github.com/testcontainers/testcontainers-go/wait"
+)
+
+func ExampleExecStrategy() {
+	ctx := context.Background()
+	req := testcontainers.ContainerRequest{
+		Image:      "localstack/localstack:latest",
+		WaitingFor: wait.ForExec([]string{"awslocal", "dynamodb", "list-tables"}),
+	}
+
+	localstack, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
+		ContainerRequest: req,
+		Started:          true,
+	})
+	if err != nil {
+		panic(err)
+	}
+
+	defer localstack.Terminate(ctx) // nolint: errcheck
+	// Here you have a running container
+}
+
+type mockExecTarget struct {
+	waitDuration time.Duration
+	successAfter time.Time
+	exitCode     int
+	failure      error
+}
+
+func (st mockExecTarget) Host(_ context.Context) (string, error) {
+	return "", errors.New("not implemented")
+}
+
+func (st mockExecTarget) MappedPort(_ context.Context, n nat.Port) (nat.Port, error) {
+	return n, errors.New("not implemented")
+}
+
+func (st mockExecTarget) Logs(_ context.Context) (io.ReadCloser, error) {
+	return nil, errors.New("not implemented")
+}
+
+func (st mockExecTarget) Exec(ctx context.Context, _ []string) (int, error) {
+	time.Sleep(st.waitDuration)
+
+	if err := ctx.Err(); err != nil {
+		return st.exitCode, err
+	}
+
+	if !st.successAfter.IsZero() && time.Now().After(st.successAfter) {
+		return 0, st.failure
+	}
+
+	return st.exitCode, st.failure
+}
+
+func (st mockExecTarget) State(_ context.Context) (*types.ContainerState, error) {
+	return nil, errors.New("not implemented")
+}
+
+func TestExecStrategyWaitUntilReady(t *testing.T) {
+	target := mockExecTarget{}
+	wg := wait.NewExecStrategy([]string{"true"}).
+		WithStartupTimeout(30 * time.Second)
+	err := wg.WaitUntilReady(context.Background(), target)
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestExecStrategyWaitUntilReadyForExec(t *testing.T) {
+	target := mockExecTarget{}
+	wg := wait.ForExec([]string{"true"})
+	err := wg.WaitUntilReady(context.Background(), target)
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestExecStrategyWaitUntilReady_MultipleChecks(t *testing.T) {
+	target := mockExecTarget{
+		exitCode:     10,
+		successAfter: time.Now().Add(2 * time.Second),
+	}
+	wg := wait.NewExecStrategy([]string{"true"}).
+		WithPollInterval(500 * time.Millisecond)
+	err := wg.WaitUntilReady(context.Background(), target)
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestExecStrategyWaitUntilReady_DeadlineExceeded(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
+	defer cancel()
+
+	target := mockExecTarget{
+		waitDuration: 1 * time.Second,
+	}
+	wg := wait.NewExecStrategy([]string{"true"})
+	err := wg.WaitUntilReady(ctx, target)
+	if err != context.DeadlineExceeded {
+		t.Fatal(err)
+	}
+}
+
+func TestExecStrategyWaitUntilReady_CustomExitCode(t *testing.T) {
+	target := mockExecTarget{
+		exitCode: 10,
+	}
+	wg := wait.NewExecStrategy([]string{"true"}).WithExitCodeMatcher(func(exitCode int) bool {
+		return exitCode == 10
+	})
+	err := wg.WaitUntilReady(context.Background(), target)
+	if err != nil {
+		t.Fatal(err)
+	}
+}