Skip to content

Commit f111b08

Browse files
authored
Added process stubbing for easier testing of launched subprocesses (#963)
## Changes This PR makes unit testing with subprocesses fast. ``` 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")) ``` This should make further iterations of #914 easier ## Tests `make test`
1 parent d4c0027 commit f111b08

File tree

4 files changed

+237
-7
lines changed

4 files changed

+237
-7
lines changed

libs/process/background.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ func Background(ctx context.Context, args []string, opts ...execOption) (string,
4747
return "", err
4848
}
4949
}
50-
if err := cmd.Run(); err != nil {
50+
if err := runCmd(ctx, cmd); err != nil {
5151
return stdout.String(), &ProcessError{
5252
Err: err,
5353
Command: commandStr,

libs/process/forwarded.go

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,10 +34,5 @@ func Forwarded(ctx context.Context, args []string, src io.Reader, outWriter, err
3434
}
3535
}
3636

37-
err := cmd.Start()
38-
if err != nil {
39-
return err
40-
}
41-
42-
return cmd.Wait()
37+
return runCmd(ctx, cmd)
4338
}

libs/process/stub.go

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
package process
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"os/exec"
7+
"path/filepath"
8+
"strings"
9+
)
10+
11+
var stubKey int
12+
13+
// WithStub creates process stub for fast and flexible testing of subprocesses
14+
func WithStub(ctx context.Context) (context.Context, *processStub) {
15+
stub := &processStub{responses: map[string]reponseStub{}}
16+
ctx = context.WithValue(ctx, &stubKey, stub)
17+
return ctx, stub
18+
}
19+
20+
func runCmd(ctx context.Context, cmd *exec.Cmd) error {
21+
stub, ok := ctx.Value(&stubKey).(*processStub)
22+
if ok {
23+
return stub.run(cmd)
24+
}
25+
return cmd.Run()
26+
}
27+
28+
type reponseStub struct {
29+
stdout string
30+
stderr string
31+
err error
32+
}
33+
34+
type processStub struct {
35+
reponseStub
36+
calls []*exec.Cmd
37+
callback func(*exec.Cmd) error
38+
responses map[string]reponseStub
39+
}
40+
41+
func (s *processStub) WithStdout(output string) *processStub {
42+
s.reponseStub.stdout = output
43+
return s
44+
}
45+
46+
func (s *processStub) WithFailure(err error) *processStub {
47+
s.reponseStub.err = err
48+
return s
49+
}
50+
51+
func (s *processStub) WithCallback(cb func(cmd *exec.Cmd) error) *processStub {
52+
s.callback = cb
53+
return s
54+
}
55+
56+
// WithStdoutFor predefines standard output response for a command. The first word
57+
// in the command string is the executable name, and NOT the executable path.
58+
// The following command would stub "2" output for "/usr/local/bin/echo 1" command:
59+
//
60+
// stub.WithStdoutFor("echo 1", "2")
61+
func (s *processStub) WithStdoutFor(command, out string) *processStub {
62+
s.responses[command] = reponseStub{
63+
stdout: out,
64+
stderr: s.responses[command].stderr,
65+
err: s.responses[command].err,
66+
}
67+
return s
68+
}
69+
70+
// WithStderrFor same as [WithStdoutFor], but for standard error
71+
func (s *processStub) WithStderrFor(command, out string) *processStub {
72+
s.responses[command] = reponseStub{
73+
stderr: out,
74+
stdout: s.responses[command].stdout,
75+
err: s.responses[command].err,
76+
}
77+
return s
78+
}
79+
80+
// WithFailureFor same as [WithStdoutFor], but for process failures
81+
func (s *processStub) WithFailureFor(command string, err error) *processStub {
82+
s.responses[command] = reponseStub{
83+
err: err,
84+
stderr: s.responses[command].stderr,
85+
stdout: s.responses[command].stdout,
86+
}
87+
return s
88+
}
89+
90+
func (s *processStub) String() string {
91+
return fmt.Sprintf("process stub with %d calls", s.Len())
92+
}
93+
94+
func (s *processStub) Len() int {
95+
return len(s.calls)
96+
}
97+
98+
func (s *processStub) Commands() (called []string) {
99+
for _, v := range s.calls {
100+
called = append(called, s.normCmd(v))
101+
}
102+
return
103+
}
104+
105+
// CombinedEnvironment returns all enviroment variables used for all commands
106+
func (s *processStub) CombinedEnvironment() map[string]string {
107+
environment := map[string]string{}
108+
for _, cmd := range s.calls {
109+
for _, line := range cmd.Env {
110+
k, v, ok := strings.Cut(line, "=")
111+
if !ok {
112+
continue
113+
}
114+
environment[k] = v
115+
}
116+
}
117+
return environment
118+
}
119+
120+
// LookupEnv returns a value from any of the triggered process environments
121+
func (s *processStub) LookupEnv(key string) string {
122+
environment := s.CombinedEnvironment()
123+
return environment[key]
124+
}
125+
126+
func (s *processStub) normCmd(v *exec.Cmd) string {
127+
// to reduce testing noise, we collect here only the deterministic binary basenames, e.g.
128+
// "/var/folders/bc/7qf8yghj6v14t40096pdcqy40000gp/T/tmp.03CAcYcbOI/python3" becomes "python3".
129+
// Use [processStub.WithCallback] if you need to match against the full executable path.
130+
binaryName := filepath.Base(v.Path)
131+
args := strings.Join(v.Args[1:], " ")
132+
return fmt.Sprintf("%s %s", binaryName, args)
133+
}
134+
135+
func (s *processStub) run(cmd *exec.Cmd) error {
136+
s.calls = append(s.calls, cmd)
137+
resp, ok := s.responses[s.normCmd(cmd)]
138+
if ok {
139+
if resp.stdout != "" {
140+
cmd.Stdout.Write([]byte(resp.stdout))
141+
}
142+
if resp.stderr != "" {
143+
cmd.Stderr.Write([]byte(resp.stderr))
144+
}
145+
return resp.err
146+
}
147+
if s.callback != nil {
148+
return s.callback(cmd)
149+
}
150+
if s.reponseStub.stdout != "" {
151+
cmd.Stdout.Write([]byte(s.reponseStub.stdout))
152+
}
153+
return s.reponseStub.err
154+
}

libs/process/stub_test.go

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
package process_test
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"os/exec"
7+
"testing"
8+
9+
"github.com/databricks/cli/libs/env"
10+
"github.com/databricks/cli/libs/process"
11+
"github.com/stretchr/testify/require"
12+
)
13+
14+
func TestStubOutput(t *testing.T) {
15+
ctx := context.Background()
16+
ctx, stub := process.WithStub(ctx)
17+
stub.WithStdout("meeee")
18+
19+
ctx = env.Set(ctx, "FOO", "bar")
20+
21+
out, err := process.Background(ctx, []string{"/usr/local/bin/meeecho", "1", "--foo", "bar"})
22+
require.NoError(t, err)
23+
require.Equal(t, "meeee", out)
24+
require.Equal(t, 1, stub.Len())
25+
require.Equal(t, []string{"meeecho 1 --foo bar"}, stub.Commands())
26+
27+
allEnv := stub.CombinedEnvironment()
28+
require.Equal(t, "bar", allEnv["FOO"])
29+
require.Equal(t, "bar", stub.LookupEnv("FOO"))
30+
}
31+
32+
func TestStubFailure(t *testing.T) {
33+
ctx := context.Background()
34+
ctx, stub := process.WithStub(ctx)
35+
stub.WithFailure(fmt.Errorf("nope"))
36+
37+
_, err := process.Background(ctx, []string{"/bin/meeecho", "1"})
38+
require.EqualError(t, err, "/bin/meeecho 1: nope")
39+
require.Equal(t, 1, stub.Len())
40+
}
41+
42+
func TestStubCallback(t *testing.T) {
43+
ctx := context.Background()
44+
ctx, stub := process.WithStub(ctx)
45+
stub.WithCallback(func(cmd *exec.Cmd) error {
46+
cmd.Stderr.Write([]byte("something..."))
47+
cmd.Stdout.Write([]byte("else..."))
48+
return fmt.Errorf("yep")
49+
})
50+
51+
_, err := process.Background(ctx, []string{"/bin/meeecho", "1"})
52+
require.EqualError(t, err, "/bin/meeecho 1: yep")
53+
require.Equal(t, 1, stub.Len())
54+
55+
var processError *process.ProcessError
56+
require.ErrorAs(t, err, &processError)
57+
require.Equal(t, "something...", processError.Stderr)
58+
require.Equal(t, "else...", processError.Stdout)
59+
}
60+
61+
func TestStubResponses(t *testing.T) {
62+
ctx := context.Background()
63+
ctx, stub := process.WithStub(ctx)
64+
stub.
65+
WithStdoutFor("qux 1", "first").
66+
WithStdoutFor("qux 2", "second").
67+
WithFailureFor("qux 3", fmt.Errorf("nope"))
68+
69+
first, err := process.Background(ctx, []string{"/path/is/irrelevant/qux", "1"})
70+
require.NoError(t, err)
71+
require.Equal(t, "first", first)
72+
73+
second, err := process.Background(ctx, []string{"/path/is/irrelevant/qux", "2"})
74+
require.NoError(t, err)
75+
require.Equal(t, "second", second)
76+
77+
_, err = process.Background(ctx, []string{"/path/is/irrelevant/qux", "3"})
78+
require.EqualError(t, err, "/path/is/irrelevant/qux 3: nope")
79+
80+
require.Equal(t, "process stub with 3 calls", stub.String())
81+
}

0 commit comments

Comments
 (0)