diff --git a/libs/process/background.go b/libs/process/background.go index 26178a1dcf..2649d0ef2c 100644 --- a/libs/process/background.go +++ b/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, diff --git a/libs/process/forwarded.go b/libs/process/forwarded.go index df3c2dbd7d..1d7fdb71e4 100644 --- a/libs/process/forwarded.go +++ b/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) } diff --git a/libs/process/stub.go b/libs/process/stub.go new file mode 100644 index 0000000000..280a9a8a24 --- /dev/null +++ b/libs/process/stub.go @@ -0,0 +1,154 @@ +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 + callback func(*exec.Cmd) error + responses map[string]reponseStub +} + +func (s *processStub) WithStdout(output string) *processStub { + s.reponseStub.stdout = output + return s +} + +func (s *processStub) WithFailure(err error) *processStub { + s.reponseStub.err = err + return s +} + +func (s *processStub) WithCallback(cb func(cmd *exec.Cmd) error) *processStub { + s.callback = 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, + stderr: s.responses[command].stderr, + err: s.responses[command].err, + } + return s +} + +// WithStderrFor same as [WithStdoutFor], but for standard error +func (s *processStub) WithStderrFor(command, out string) *processStub { + s.responses[command] = reponseStub{ + stderr: out, + stdout: s.responses[command].stdout, + err: s.responses[command].err, + } + 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, + stderr: s.responses[command].stderr, + stdout: s.responses[command].stdout, + } + 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". + // Use [processStub.WithCallback] if you need to match against the full executable path. + 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.callback != nil { + return s.callback(cmd) + } + if s.reponseStub.stdout != "" { + cmd.Stdout.Write([]byte(s.reponseStub.stdout)) + } + return s.reponseStub.err +} diff --git a/libs/process/stub_test.go b/libs/process/stub_test.go new file mode 100644 index 0000000000..65f59f8172 --- /dev/null +++ b/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.WithStdout("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.WithFailure(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.WithCallback(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()) +}