| 1 | package harness |
| 2 | |
| 3 | import ( |
| 4 | "fmt" |
| 5 | "io" |
| 6 | "os" |
| 7 | "os/exec" |
| 8 | "strings" |
| 9 | ) |
| 10 | |
| 11 | // Runner is a process runner which can run subprocesses and aggregate output. |
| 12 | type Runner struct { |
| 13 | Env map[string]string |
| 14 | Dir string |
| 15 | Verbose bool |
| 16 | } |
| 17 | |
| 18 | type ( |
| 19 | CmdOpt func(*exec.Cmd) |
| 20 | RunFunc func(*exec.Cmd) error |
| 21 | ) |
| 22 | |
| 23 | var RunFuncStart = (*exec.Cmd).Start |
| 24 | |
| 25 | type RunRequest struct { |
| 26 | Path string |
| 27 | Args []string |
| 28 | // Options that are applied to the exec.Cmd just before running it |
| 29 | CmdOpts []CmdOpt |
| 30 | // Function to use to run the command. |
| 31 | // If not specified, defaults to cmd.Run |
| 32 | RunFunc func(*exec.Cmd) error |
| 33 | Verbose bool |
| 34 | } |
| 35 | |
| 36 | type RunResult struct { |
| 37 | Stdout *Buffer |
| 38 | Stderr *Buffer |
| 39 | Err error |
| 40 | ExitErr *exec.ExitError |
| 41 | Cmd *exec.Cmd |
| 42 | } |
| 43 | |
| 44 | func (r *RunResult) ExitCode() int { |
| 45 | return r.Cmd.ProcessState.ExitCode() |
| 46 | } |
| 47 | |
| 48 | func environToMap(environ []string) map[string]string { |
| 49 | m := map[string]string{} |
| 50 | for _, e := range environ { |
| 51 | kv := strings.Split(e, "=") |
| 52 | // Skip environment variables that start with = |
| 53 | // These can occur in Windows https://github.com/golang/go/issues/61956 |
| 54 | if kv[0] == "" { |
| 55 | continue |
| 56 | } |
| 57 | m[kv[0]] = kv[1] |
| 58 | } |
| 59 | return m |
| 60 | } |
| 61 | |
| 62 | func (r *Runner) Run(req RunRequest) *RunResult { |
| 63 | cmd := exec.Command(req.Path, req.Args...) |
| 64 | var stdout io.Writer |
| 65 | var stderr io.Writer |
| 66 | outbuf := &Buffer{} |
| 67 | errbuf := &Buffer{} |
| 68 | |
| 69 | if r.Verbose { |
| 70 | or, ow := io.Pipe() |
| 71 | errr, errw := io.Pipe() |
| 72 | stdout = io.MultiWriter(outbuf, ow) |
| 73 | stderr = io.MultiWriter(errbuf, errw) |
| 74 | go func() { |
| 75 | _, _ = io.Copy(os.Stdout, or) |
| 76 | }() |
| 77 | go func() { |
| 78 | _, _ = io.Copy(os.Stderr, errr) |
| 79 | }() |
| 80 | } else { |
| 81 | stdout = outbuf |
| 82 | stderr = errbuf |
| 83 | } |
| 84 | |
| 85 | cmd.Stdout = stdout |
| 86 | cmd.Stderr = stderr |
| 87 | cmd.Dir = r.Dir |
| 88 | |
| 89 | for k, v := range r.Env { |
| 90 | cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", k, v)) |
| 91 | } |
| 92 | |
| 93 | for _, o := range req.CmdOpts { |
| 94 | o(cmd) |
| 95 | } |
| 96 | |
| 97 | if req.RunFunc == nil { |
| 98 | req.RunFunc = (*exec.Cmd).Run |
| 99 | } |
| 100 | |
| 101 | log.Debugf("running %v", cmd.Args) |
| 102 | |
| 103 | err := req.RunFunc(cmd) |
| 104 | |
| 105 | result := RunResult{ |
| 106 | Stdout: outbuf, |
| 107 | Stderr: errbuf, |
| 108 | Cmd: cmd, |
| 109 | Err: err, |
| 110 | } |
| 111 | |
| 112 | if exitErr, ok := err.(*exec.ExitError); ok { |
| 113 | result.ExitErr = exitErr |
| 114 | } |
| 115 | |
| 116 | return &result |
| 117 | } |
| 118 | |
| 119 | // MustRun runs the command and fails the test if the command fails. |
| 120 | func (r *Runner) MustRun(req RunRequest) *RunResult { |
| 121 | result := r.Run(req) |
| 122 | r.AssertNoError(result) |
| 123 | return result |
| 124 | } |
| 125 | |
| 126 | func (r *Runner) AssertNoError(result *RunResult) { |
| 127 | if result.ExitErr != nil { |
| 128 | log.Panicf("'%s' returned error, code: %d, err: %s\nstdout:%s\nstderr:%s\n", |
| 129 | result.Cmd.Args, result.ExitErr.ExitCode(), result.ExitErr.Error(), result.Stdout.String(), result.Stderr.String()) |
| 130 | } |
| 131 | if result.Err != nil { |
| 132 | log.Panicf("unable to run %s: %s", result.Cmd.Path, result.Err) |
| 133 | } |
| 134 | } |
| 135 | |
| 136 | func RunWithEnv(env map[string]string) CmdOpt { |
| 137 | return func(cmd *exec.Cmd) { |
| 138 | for k, v := range env { |
| 139 | cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", k, v)) |
| 140 | } |
| 141 | } |
| 142 | } |
| 143 | |
| 144 | func RunWithPath(path string) CmdOpt { |
| 145 | return func(cmd *exec.Cmd) { |
| 146 | var newEnv []string |
| 147 | for _, env := range cmd.Env { |
| 148 | e := strings.Split(env, "=") |
| 149 | if e[0] == "PATH" { |
| 150 | paths := strings.Split(e[1], ":") |
| 151 | paths = append(paths, path) |
| 152 | e[1] = strings.Join(paths, ":") |
| 153 | fmt.Printf("path: %s\n", strings.Join(e, "=")) |
| 154 | } |
| 155 | newEnv = append(newEnv, strings.Join(e, "=")) |
| 156 | } |
| 157 | cmd.Env = newEnv |
| 158 | } |
| 159 | } |
| 160 | |
| 161 | func RunWithStdin(reader io.Reader) CmdOpt { |
| 162 | return func(cmd *exec.Cmd) { |
| 163 | cmd.Stdin = reader |
| 164 | } |
| 165 | } |
| 166 | |
| 167 | func RunWithStdinStr(s string) CmdOpt { |
| 168 | return RunWithStdin(strings.NewReader(s)) |
| 169 | } |
| 170 | |
| 171 | func RunWithStdout(writer io.Writer) CmdOpt { |
| 172 | return func(cmd *exec.Cmd) { |
| 173 | cmd.Stdout = io.MultiWriter(writer, cmd.Stdout) |
| 174 | } |
| 175 | } |
| 176 | |
| 177 | func RunWithStderr(writer io.Writer) CmdOpt { |
| 178 | return func(cmd *exec.Cmd) { |
| 179 | cmd.Stderr = io.MultiWriter(writer, cmd.Stdout) |
| 180 | } |
| 181 | } |