Use WaitDelay in non_windows script timeout (#22912)

orbit/pkg/scripts/exec_nonwindows_test.go

@@ -4,12 +4,14 @@ package scripts
 
 import (
 	"context"
+	"io/ioutil"
 	"os"
 	"os/exec"
 	"path/filepath"
 	"runtime"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/fleetdm/fleet/v4/server/fleet"
 	"github.com/stretchr/testify/require"
@@ -78,3 +80,76 @@ func TestExecCmdNonWindows(t *testing.T) {
 		})
 	}
 }
+
+func writeTestScript(content string) (string, error) {
+	tmpfile, err := ioutil.TempFile("", "testscript*.sh")
+	if err != nil {
+		return "", err
+	}
+
+	if _, err := tmpfile.Write([]byte(content)); err != nil {
+		tmpfile.Close()
+		return "", err
+	}
+	if err := tmpfile.Close(); err != nil {
+		return "", err
+	}
+
+	err = os.Chmod(tmpfile.Name(), 0o700)
+	if err != nil {
+		return "", err
+	}
+
+	return tmpfile.Name(), nil
+}
+
+func TestExecCmdTimeout(t *testing.T) {
+	scriptContent := `#!/bin/sh
+	sleep 5
+	echo "Finished"`
+	scriptPath, err := writeTestScript(scriptContent)
+	if err != nil {
+		t.Fatalf("Failed to write test script: %v", err)
+	}
+	defer os.Remove(scriptPath)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+
+	start := time.Now()
+	output, exitCode, err := ExecCmd(ctx, scriptPath, nil)
+	require.NotNil(t, err)
+	require.Contains(t, err.Error(), "signal: killed")
+	if exitCode != -1 {
+		t.Fatalf("Expected exit code -1, got: %d", exitCode)
+	}
+	if len(output) != 0 {
+		t.Fatalf("Expected no output, got: %s", output)
+	}
+	require.True(t, time.Since(start) <= 5*time.Second)
+}
+
+func TestExecCmdSuccess(t *testing.T) {
+	scriptContent := `#!/bin/sh
+	echo "Hello, World!"`
+	scriptPath, err := writeTestScript(scriptContent)
+	if err != nil {
+		t.Fatalf("Failed to write test script: %v", err)
+	}
+	defer os.Remove(scriptPath)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	output, exitCode, err := ExecCmd(ctx, scriptPath, nil)
+	if err != nil {
+		t.Fatalf("Expected no error, got: %v", err)
+	}
+	if exitCode != 0 {
+		t.Fatalf("Expected exit code 0, got: %d", exitCode)
+	}
+	expectedOutput := "Hello, World!\n"
+	if string(output) != expectedOutput {
+		t.Fatalf("Expected output %q, got: %q", expectedOutput, output)
+	}
+}