From 1cb6f223f775ff7baf9c8a5772205939acf367dd Mon Sep 17 00:00:00 2001
From: Barrett Strausser <bearrito@users.noreply.github.com>
Date: Wed, 27 Mar 2024 11:57:25 -0400
Subject: [PATCH] feat(exitcode): Add exit code sugar method (#2342)

* Add exit code sugar method

* Fix lint errors

* Tweak test name

* Rename test

---------

Co-authored-by: bstrausser <bstrausser@locusrobotics.com>
---
 wait/exec.go      |  6 ++++++
 wait/exec_test.go | 21 +++++++++++++++++++++
 2 files changed, 27 insertions(+)

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{