feat(exitcode): Add exit code sugar method (testcontainers#2342)

* Add exit code sugar method * Fix lint errors * Tweak test name * Rename test --------- Co-authored-by: bstrausser <[email protected]>
gustavosbarreto · Mar 27, 2024 · 1cb6f22 · 1cb6f22
1 parent 32f372c
commit 1cb6f22
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 0 deletions.
diff --git 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
@@ -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{