diff --git a/wait/exec.go b/wait/exec.go index e0ffc89758..2e341dd3e0 100644 --- a/wait/exec.go +++ b/wait/exec.go @@ -45,6 +45,12 @@ func (ws *ExecStrategy) WithStartupTimeout(startupTimeout time.Duration) *ExecSt return ws } +func (ws *ExecStrategy) WithExitCode(exitCode int) *ExecStrategy { + return ws.WithExitCodeMatcher(func(actualCode int) bool { + return actualCode == exitCode + }) +} + func (ws *ExecStrategy) WithExitCodeMatcher(exitCodeMatcher func(exitCode int) bool) *ExecStrategy { ws.ExitCodeMatcher = exitCodeMatcher return ws diff --git a/wait/exec_test.go b/wait/exec_test.go index b1c31f1b6e..132933018f 100644 --- a/wait/exec_test.go +++ b/wait/exec_test.go @@ -156,6 +156,27 @@ func TestExecStrategyWaitUntilReady_CustomExitCode(t *testing.T) { } } +func TestExecStrategyWaitUntilReady_withExitCode(t *testing.T) { + target := mockExecTarget{ + exitCode: 10, + } + wg := wait.NewExecStrategy([]string{"true"}).WithExitCode(10) + // Default is 60. Let's shorten that + wg.WithStartupTimeout(time.Second * 2) + err := wg.WaitUntilReady(context.Background(), target) + if err != nil { + t.Fatal(err) + } + + // Ensure we aren't spuriously returning on any code + wg = wait.NewExecStrategy([]string{"true"}).WithExitCode(0) + wg.WithStartupTimeout(time.Second * 2) + err = wg.WaitUntilReady(context.Background(), target) + if err == nil { + t.Fatalf("Expected strategy to timeout out") + } +} + func TestExecStrategyWaitUntilReady_CustomResponseMatcher(t *testing.T) { // waitForExecExitCodeResponse { dockerReq := testcontainers.ContainerRequest{