feat: add status check to wait strategies (#944)

wait/health.go

@@ -2,8 +2,9 @@ package wait
 
 import (
 	"context"
-	"github.com/docker/docker/api/types"
 	"time"
+
+	"github.com/docker/docker/api/types"
 )
 
 // Implement interface
@@ -77,6 +78,9 @@ func (ws *HealthStrategy) WaitUntilReady(ctx context.Context, target StrategyTar
 			if err != nil {
 				return err
 			}
+			if err := checkState(state); err != nil {
+				return err
+			}
 			if state.Health == nil || state.Health.Status != types.Healthy {
 				time.Sleep(ws.PollInterval)
 				continue

wait/health_test.go

@@ -3,18 +3,18 @@ package wait
 import (
 	"context"
 	"errors"
-	"github.com/stretchr/testify/assert"
 	"io"
 	"testing"
 	"time"
 
 	"github.com/docker/docker/api/types"
 	"github.com/docker/go-connections/nat"
+	"github.com/stretchr/testify/assert"
 	tcexec "github.com/testcontainers/testcontainers-go/exec"
 )
 
 type healthStrategyTarget struct {
-	Health *types.Health
+	state *types.ContainerState
 }
 
 func (st healthStrategyTarget) Host(ctx context.Context) (string, error) {
@@ -38,13 +38,18 @@ func (st healthStrategyTarget) Exec(ctx context.Context, cmd []string, options .
 }
 
 func (st healthStrategyTarget) State(ctx context.Context) (*types.ContainerState, error) {
-	return &types.ContainerState{Health: st.Health}, nil
+	return st.state, nil
 }
 
 // TestWaitForHealthTimesOutForUnhealthy confirms that an unhealthy container will eventually
 // time out.
 func TestWaitForHealthTimesOutForUnhealthy(t *testing.T) {
-	target := healthStrategyTarget{Health: &types.Health{Status: types.Unhealthy}}
+	target := healthStrategyTarget{
+		state: &types.ContainerState{
+			Running: true,
+			Health:  &types.Health{Status: types.Unhealthy},
+		},
+	}
 	wg := NewHealthStrategy().WithStartupTimeout(100 * time.Millisecond)
 	err := wg.WaitUntilReady(context.Background(), target)
 
@@ -54,7 +59,12 @@ func TestWaitForHealthTimesOutForUnhealthy(t *testing.T) {
 
 // TestWaitForHealthSucceeds ensures that a healthy container always succeeds.
 func TestWaitForHealthSucceeds(t *testing.T) {
-	target := healthStrategyTarget{Health: &types.Health{Status: types.Healthy}}
+	target := healthStrategyTarget{
+		state: &types.ContainerState{
+			Running: true,
+			Health:  &types.Health{Status: types.Healthy},
+		},
+	}
 	wg := NewHealthStrategy().WithStartupTimeout(100 * time.Millisecond)
 	err := wg.WaitUntilReady(context.Background(), target)
 
@@ -64,7 +74,12 @@ func TestWaitForHealthSucceeds(t *testing.T) {
 // TestWaitForHealthWithNil checks that an initial `nil` Health will not casue a panic,
 // and if the container eventually becomes healthy, the HealthStrategy will succeed.
 func TestWaitForHealthWithNil(t *testing.T) {
-	target := &healthStrategyTarget{Health: nil}
+	target := &healthStrategyTarget{
+		state: &types.ContainerState{
+			Running: true,
+			Health:  nil,
+		},
+	}
 	wg := NewHealthStrategy().
 		WithStartupTimeout(500 * time.Millisecond).
 		WithPollInterval(100 * time.Millisecond)
@@ -73,7 +88,7 @@ func TestWaitForHealthWithNil(t *testing.T) {
 		// wait a bit to simulate startup time and give check time to at least
 		// try a few times with a nil Health
 		time.Sleep(200 * time.Millisecond)
-		target.Health = &types.Health{Status: types.Healthy}
+		target.state.Health = &types.Health{Status: types.Healthy}
 	}(target)
 
 	err := wg.WaitUntilReady(context.Background(), target)
@@ -82,7 +97,12 @@ func TestWaitForHealthWithNil(t *testing.T) {
 
 // TestWaitFailsForNilHealth checks that Health always nil fails (but will NOT cause a panic)
 func TestWaitFailsForNilHealth(t *testing.T) {
-	target := &healthStrategyTarget{Health: nil}
+	target := &healthStrategyTarget{
+		state: &types.ContainerState{
+			Running: true,
+			Health:  nil,
+		},
+	}
 	wg := NewHealthStrategy().
 		WithStartupTimeout(500 * time.Millisecond).
 		WithPollInterval(100 * time.Millisecond)
@@ -91,3 +111,49 @@ func TestWaitFailsForNilHealth(t *testing.T) {
 	assert.NotNil(t, err)
 	assert.True(t, errors.Is(err, context.DeadlineExceeded))
 }
+
+func TestWaitForHealthFailsDueToOOMKilledContainer(t *testing.T) {
+	target := &healthStrategyTarget{
+		state: &types.ContainerState{
+			OOMKilled: true,
+		},
+	}
+	wg := NewHealthStrategy().
+		WithStartupTimeout(500 * time.Millisecond).
+		WithPollInterval(100 * time.Millisecond)
+
+	err := wg.WaitUntilReady(context.Background(), target)
+	assert.NotNil(t, err)
+	assert.EqualError(t, err, "container crashed with out-of-memory (OOMKilled)")
+}
+
+func TestWaitForHealthFailsDueToExitedContainer(t *testing.T) {
+	target := &healthStrategyTarget{
+		state: &types.ContainerState{
+			Status:   "exited",
+			ExitCode: 1,
+		},
+	}
+	wg := NewHealthStrategy().
+		WithStartupTimeout(500 * time.Millisecond).
+		WithPollInterval(100 * time.Millisecond)
+
+	err := wg.WaitUntilReady(context.Background(), target)
+	assert.NotNil(t, err)
+	assert.EqualError(t, err, "container exited with code 1")
+}
+
+func TestWaitForHealthFailsDueToUnexpectedContainerStatus(t *testing.T) {
+	target := &healthStrategyTarget{
+		state: &types.ContainerState{
+			Status: "dead",
+		},
+	}
+	wg := NewHealthStrategy().
+		WithStartupTimeout(500 * time.Millisecond).
+		WithPollInterval(100 * time.Millisecond)
+
+	err := wg.WaitUntilReady(context.Background(), target)
+	assert.NotNil(t, err)
+	assert.EqualError(t, err, "unexpected container status \"dead\"")
+}

wait/host_port.go

@@ -113,6 +113,9 @@ func (hp *HostPortStrategy) WaitUntilReady(ctx context.Context, target StrategyT
 		case <-ctx.Done():
 			return fmt.Errorf("%s:%w", ctx.Err(), err)
 		case <-time.After(waitInterval):
+			if err := checkTarget(ctx, target); err != nil {
+				return err
+			}
 			port, err = target.MappedPort(ctx, internalPort)
 			if err != nil {
 				fmt.Printf("(%d) [%s] %s\n", i, port, err)
@@ -128,6 +131,9 @@ func (hp *HostPortStrategy) WaitUntilReady(ctx context.Context, target StrategyT
 	dialer := net.Dialer{}
 	address := net.JoinHostPort(ipAddress, portString)
 	for {
+		if err := checkTarget(ctx, target); err != nil {
+			return err
+		}
 		conn, err := dialer.DialContext(ctx, proto, address)
 		if err != nil {
 			if v, ok := err.(*net.OpError); ok {
@@ -151,6 +157,9 @@ func (hp *HostPortStrategy) WaitUntilReady(ctx context.Context, target StrategyT
 		if ctx.Err() != nil {
 			return ctx.Err()
 		}
+		if err := checkTarget(ctx, target); err != nil {
+			return err
+		}
 		exitCode, _, err := target.Exec(ctx, []string{"/bin/sh", "-c", command})
 		if err != nil {
 			return fmt.Errorf("%w, host port waiting failed", err)