1 files changed, 78 insertions, 2 deletions

diff --git a/integrationtests/commandutils.go b/integrationtests/commandutils.go
index 8d81955..1947de4 100644
--- a/integrationtests/commandutils.go
+++ b/integrationtests/commandutils.go
@@ -2,6 +2,7 @@ package integrationtests
 
 import (
 	"bufio"
+	"bytes"
 	"context"
 	"fmt"
 	"io"
@@ -15,23 +16,38 @@ import (
 
 func runCommand(ctx context.Context, t *testing.T, stdoutFile, cmdStr string,
 	args ...string) (int, error) {
+	return runCommandWithEnv(ctx, t, stdoutFile, cmdStr, nil, args...)
+}
 
+func runCommandWithEnv(ctx context.Context, t *testing.T, stdoutFile, cmdStr string,
+	env map[string]string, args ...string) (int, error) {
 	if _, err := os.Stat(cmdStr); err != nil {
 		return 0, fmt.Errorf("no such executable '%s', please compile first: %w", cmdStr, err)
 	}
 
+	// Log command execution if logger is available
+	if logger := GetTestLogger(ctx); logger != nil {
+		logger.LogCommand(cmdStr, args)
+	}
+
 	t.Log("Creating stdout file", stdoutFile)
 	fd, err := os.Create(stdoutFile)
 	if err != nil {
-		return 0, nil
+		return 0, err
 	}
 	defer fd.Close()
 
 	t.Log("Running command", cmdStr, strings.Join(args, " "))
 	cmd := exec.CommandContext(ctx, cmdStr, args...)
+	cmd.Env = os.Environ()
+	for key, value := range env {
+		cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", key, value))
+	}
 	out, err := cmd.CombinedOutput()
 	t.Log("Done running command!", err)
-	_, _ = fd.Write(out)
+	if _, copyErr := io.Copy(fd, bytes.NewReader(out)); copyErr != nil {
+		return exitCodeFromError(err), copyErr
+	}
 
 	return exitCodeFromError(err), err
 }
@@ -48,8 +64,54 @@ func runCommandRetry(ctx context.Context, t *testing.T, retries int, stdoutFile,
 	return
 }
 
+func runCommandUntilValid(ctx context.Context, t *testing.T, attempts int, delay time.Duration,
+	stdoutFile, cmd string, validate func() error, args ...string) error {
+
+	t.Helper()
+
+	if attempts < 1 {
+		attempts = 1
+	}
+
+	var lastErr error
+	for i := 0; i < attempts; i++ {
+		exitCode, err := runCommand(ctx, t, stdoutFile, cmd, args...)
+		if err == nil {
+			validateErr := validate()
+			if validateErr == nil {
+				return nil
+			}
+			lastErr = validateErr
+		} else {
+			lastErr = fmt.Errorf("command %s failed with exit code %d: %w", cmd, exitCode, err)
+		}
+
+		if i == attempts-1 {
+			break
+		}
+
+		timer := time.NewTimer(delay)
+		select {
+		case <-ctx.Done():
+			timer.Stop()
+			if lastErr != nil {
+				return lastErr
+			}
+			return ctx.Err()
+		case <-timer.C:
+		}
+	}
+
+	return lastErr
+}
+
 func startCommand(ctx context.Context, t *testing.T, inPipeFile,
 	cmdStr string, args ...string) (<-chan string, <-chan string, <-chan error, error) {
+	return startCommandWithEnv(ctx, t, inPipeFile, cmdStr, nil, args...)
+}
+
+func startCommandWithEnv(ctx context.Context, t *testing.T, inPipeFile,
+	cmdStr string, env map[string]string, args ...string) (<-chan string, <-chan string, <-chan error, error) {
 
 	stdoutCh := make(chan string)
 	stderrCh := make(chan string)
@@ -59,9 +121,23 @@ func startCommand(ctx context.Context, t *testing.T, inPipeFile,
 			fmt.Errorf("no such executable '%s', please compile first: %w", cmdStr, err)
 	}
 
+	// Log command execution if logger is available
+	if logger := GetTestLogger(ctx); logger != nil {
+		logger.LogCommand(cmdStr, args)
+	}
+
 	t.Log(cmdStr, strings.Join(args, " "))
 	cmd := exec.CommandContext(ctx, cmdStr, args...)
 
+	// Always inherit environment variables
+	cmd.Env = os.Environ()
+	// Add any additional environment variables if provided
+	if env != nil {
+		for k, v := range env {
+			cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", k, v))
+		}
+	}
+
 	var stdinPipe io.WriteCloser
 	if inPipeFile != "" {
 		var err error