Added process stubbing for easier testing of launched subprocesses

libs/process/background.go

@@ -47,7 +47,7 @@ func Background(ctx context.Context, args []string, opts ...execOption) (string,
 			return "", err
 		}
 	}
-	if err := cmd.Run(); err != nil {
+	if err := runCmd(ctx, cmd); err != nil {
 		return stdout.String(), &ProcessError{
 			Err:     err,
 			Command: commandStr,

libs/process/forwarded.go

@@ -34,10 +34,5 @@ func Forwarded(ctx context.Context, args []string, src io.Reader, outWriter, err
 		}
 	}
 
-	err := cmd.Start()
-	if err != nil {
-		return err
-	}
-
-	return cmd.Wait()
+	return runCmd(ctx, cmd)
 }

libs/process/stub.go

@@ -0,0 +1,149 @@
+package process
+
+import (
+	"context"
+	"fmt"
+	"os/exec"
+	"path/filepath"
+	"strings"
+)
+
+var stubKey int
+
+// WithStub creates process stub for fast and flexible testing of subprocesses
+func WithStub(ctx context.Context) (context.Context, *processStub) {
+	stub := &processStub{responses: map[string]reponseStub{}}
+	ctx = context.WithValue(ctx, &stubKey, stub)
+	return ctx, stub
+}
+
+func runCmd(ctx context.Context, cmd *exec.Cmd) error {
+	stub, ok := ctx.Value(&stubKey).(*processStub)
+	if ok {
+		return stub.run(cmd)
+	}
+	return cmd.Run()
+}
+
+type reponseStub struct {
+	stdout string
+	stderr string
+	err    error
+}
+
+type processStub struct {
+	reponseStub
+	calls           []*exec.Cmd
+	defaultCallback func(*exec.Cmd) error
+	responses       map[string]reponseStub
+}
+
+func (s *processStub) WithDefaultOutput(output string) *processStub {
+	s.reponseStub.stdout = output
+	return s
+}
+
+func (s *processStub) WithDefaultFailure(err error) *processStub {
+	s.reponseStub.err = err
+	return s
+}
+
+func (s *processStub) WithDefaultCallback(cb func(cmd *exec.Cmd) error) *processStub {
+	s.defaultCallback = cb
+	return s
+}
+
+// WithStdoutFor predefines standard output response for a command. The first word
+// in the command string is the executable name, and NOT the executable path.
+// The following command would stub "2" output for "/usr/local/bin/echo 1" command:
+//
+//	stub.WithStdoutFor("echo 1", "2")
+func (s *processStub) WithStdoutFor(command, out string) *processStub {
+	s.responses[command] = reponseStub{
+		stdout: out,
+	}
+	return s
+}
+
+// WithStdoutFor same as [WithStdoutFor], but for standard error
+func (s *processStub) WithStderrFor(command, out string) *processStub {
+	s.responses[command] = reponseStub{
+		stderr: out,
+	}
+	return s
+}
+
+// WithFailureFor same as [WithStdoutFor], but for process failures
+func (s *processStub) WithFailureFor(command string, err error) *processStub {
+	s.responses[command] = reponseStub{
+		err: err,
+	}
+	return s
+}
+
+func (s *processStub) String() string {
+	return fmt.Sprintf("process stub with %d calls", s.Len())
+}
+
+func (s *processStub) Len() int {
+	return len(s.calls)
+}
+
+func (s *processStub) Commands() (called []string) {
+	for _, v := range s.calls {
+		called = append(called, s.normCmd(v))
+	}
+	return
+}
+
+// CombinedEnvironment returns all enviroment variables used for all commands
+func (s *processStub) CombinedEnvironment() map[string]string {
+	environment := map[string]string{}
+	for _, cmd := range s.calls {
+		for _, line := range cmd.Env {
+			k, v, ok := strings.Cut(line, "=")
+			if !ok {
+				continue
+			}
+			environment[k] = v
+		}
+	}
+	return environment
+}
+
+// LookupEnv returns a value from any of the triggered process environments
+func (s *processStub) LookupEnv(key string) string {
+	environment := s.CombinedEnvironment()
+	return environment[key]
+}
+
+func (s *processStub) normCmd(v *exec.Cmd) string {
+	// to reduce testing noise, we collect here only the deterministic binary basenames, e.g.
+	// "/var/folders/bc/7qf8yghj6v14t40096pdcqy40000gp/T/tmp.03CAcYcbOI/python3" becomes "python3",
+	// while still giving the possibility to customize process stubbing even further.
+	// See [processStub.WithDefaultCallback]
+	binaryName := filepath.Base(v.Path)
+	args := strings.Join(v.Args[1:], " ")
+	return fmt.Sprintf("%s %s", binaryName, args)
+}
+
+func (s *processStub) run(cmd *exec.Cmd) error {
+	s.calls = append(s.calls, cmd)
+	resp, ok := s.responses[s.normCmd(cmd)]
+	if ok {
+		if resp.stdout != "" {
+			cmd.Stdout.Write([]byte(resp.stdout))
+		}
+		if resp.stderr != "" {
+			cmd.Stderr.Write([]byte(resp.stderr))
+		}
+		return resp.err
+	}
+	if s.defaultCallback != nil {
+		return s.defaultCallback(cmd)
+	}
+	if s.reponseStub.stdout != "" {
+		cmd.Stdout.Write([]byte(s.reponseStub.stdout))
+	}
+	return s.reponseStub.err
+}

libs/process/stub_test.go

@@ -0,0 +1,81 @@
+package process_test
+
+import (
+	"context"
+	"fmt"
+	"os/exec"
+	"testing"
+
+	"github.com/databricks/cli/libs/env"
+	"github.com/databricks/cli/libs/process"
+	"github.com/stretchr/testify/require"
+)
+
+func TestStubOutput(t *testing.T) {
+	ctx := context.Background()
+	ctx, stub := process.WithStub(ctx)
+	stub.WithDefaultOutput("meeee")
+
+	ctx = env.Set(ctx, "FOO", "bar")
+
+	out, err := process.Background(ctx, []string{"/usr/local/bin/meeecho", "1", "--foo", "bar"})
+	require.NoError(t, err)
+	require.Equal(t, "meeee", out)
+	require.Equal(t, 1, stub.Len())
+	require.Equal(t, []string{"meeecho 1 --foo bar"}, stub.Commands())
+
+	allEnv := stub.CombinedEnvironment()
+	require.Equal(t, "bar", allEnv["FOO"])
+	require.Equal(t, "bar", stub.LookupEnv("FOO"))
+}
+
+func TestStubFailure(t *testing.T) {
+	ctx := context.Background()
+	ctx, stub := process.WithStub(ctx)
+	stub.WithDefaultFailure(fmt.Errorf("nope"))
+
+	_, err := process.Background(ctx, []string{"/bin/meeecho", "1"})
+	require.EqualError(t, err, "/bin/meeecho 1: nope")
+	require.Equal(t, 1, stub.Len())
+}
+
+func TestStubCallback(t *testing.T) {
+	ctx := context.Background()
+	ctx, stub := process.WithStub(ctx)
+	stub.WithDefaultCallback(func(cmd *exec.Cmd) error {
+		cmd.Stderr.Write([]byte("something..."))
+		cmd.Stdout.Write([]byte("else..."))
+		return fmt.Errorf("yep")
+	})
+
+	_, err := process.Background(ctx, []string{"/bin/meeecho", "1"})
+	require.EqualError(t, err, "/bin/meeecho 1: yep")
+	require.Equal(t, 1, stub.Len())
+
+	var processError *process.ProcessError
+	require.ErrorAs(t, err, &processError)
+	require.Equal(t, "something...", processError.Stderr)
+	require.Equal(t, "else...", processError.Stdout)
+}
+
+func TestStubResponses(t *testing.T) {
+	ctx := context.Background()
+	ctx, stub := process.WithStub(ctx)
+	stub.
+		WithStdoutFor("qux 1", "first").
+		WithStdoutFor("qux 2", "second").
+		WithFailureFor("qux 3", fmt.Errorf("nope"))
+
+	first, err := process.Background(ctx, []string{"/path/is/irrelevant/qux", "1"})
+	require.NoError(t, err)
+	require.Equal(t, "first", first)
+
+	second, err := process.Background(ctx, []string{"/path/is/irrelevant/qux", "2"})
+	require.NoError(t, err)
+	require.Equal(t, "second", second)
+
+	_, err = process.Background(ctx, []string{"/path/is/irrelevant/qux", "3"})
+	require.EqualError(t, err, "/path/is/irrelevant/qux 3: nope")
+
+	require.Equal(t, "process stub with 3 calls", stub.String())
+}