databricks-cli/libs/process/stub.go

171 lines
4.2 KiB
Go
Raw Normal View History

package process
import (
"context"
"errors"
"fmt"
"os/exec"
"path/filepath"
"regexp"
"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)
var unixArgs []string
for _, arg := range v.Args[1:] {
unixArgs = append(unixArgs, filepath.ToSlash(arg))
}
args := strings.Join(unixArgs, " ")
return fmt.Sprintf("%s %s", binaryName, args)
}
var ErrStubContinue = errors.New("continue executing the stub after callback")
func (s *processStub) run(cmd *exec.Cmd) error {
s.calls = append(s.calls, cmd)
for pattern, resp := range s.responses {
re := regexp.MustCompile(pattern)
norm := s.normCmd(cmd)
if !re.MatchString(norm) {
continue
}
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)
}
var zeroStub reponseStub
if s.reponseStub == zeroStub {
return fmt.Errorf("no default process stub")
}
if s.reponseStub.stdout != "" {
cmd.Stdout.Write([]byte(s.reponseStub.stdout))
}
return s.reponseStub.err
}