Skip to content

Commit 7eb0f24

Browse files
heppuheppu
committed
Add cmd dependency
1 parent c9872d7 commit 7eb0f24

File tree

2 files changed

+365
-0
lines changed

2 files changed

+365
-0
lines changed

‎dep/cmd/cmd.go‎

Lines changed: 201 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,201 @@
1+
package cmd
2+
3+
import (
4+
"bufio"
5+
"errors"
6+
"fmt"
7+
"os"
8+
"os/exec"
9+
"reflect"
10+
"regexp"
11+
"runtime"
12+
"strings"
13+
"time"
14+
15+
"github.com/go-tstr/tstr/strerr"
16+
)
17+
18+
const (
19+
ErrMissingCmd = strerr.Error("missing command")
20+
ErrStartFailed = strerr.Error("failed to start command")
21+
ErrReadyFailed = strerr.Error("failed to verify readiness")
22+
ErrStopFailed = strerr.Error("command didn't stop successfully")
23+
ErrOptApply = strerr.Error("failed apply Opt")
24+
ErrNoMatchingLine = strerr.Error("no matching line found")
25+
ErrNilCmdRegexp = strerr.Error("command has to be set before this option can be applied, check the order of options")
26+
)
27+
28+
type Cmd struct {
29+
opts []Opt
30+
ready func(*exec.Cmd) error
31+
stop func(*exec.Cmd) error
32+
cmd *exec.Cmd
33+
readyTimeout time.Duration
34+
}
35+
36+
type Opt func(*Cmd) error
37+
38+
func New(opts ...Opt) *Cmd {
39+
return &Cmd{
40+
opts: opts,
41+
ready: func(*exec.Cmd) error { return nil },
42+
stop: StopWithSignal(os.Interrupt),
43+
readyTimeout: 30 * time.Second,
44+
}
45+
}
46+
47+
func (c *Cmd) Start() error {
48+
for _, opt := range c.opts {
49+
if err := opt(c); err != nil {
50+
return fmt.Errorf("failed to apply option %s: %w", getFnName(opt), err)
51+
}
52+
}
53+
54+
if c.cmd == nil {
55+
return ErrMissingCmd
56+
}
57+
58+
return c.wrapErr(ErrStartFailed, c.cmd.Start())
59+
}
60+
61+
func (c *Cmd) Ready() error {
62+
errCh := make(chan error, 1)
63+
go func() {
64+
defer close(errCh)
65+
errCh <- c.ready(c.cmd)
66+
}()
67+
68+
select {
69+
case <-time.After(c.readyTimeout):
70+
return c.wrapErr(ErrReadyFailed, fmt.Errorf("timeout after %s", c.readyTimeout))
71+
case err := <-errCh:
72+
return c.wrapErr(ErrReadyFailed, err)
73+
}
74+
}
75+
76+
func (c *Cmd) Stop() error {
77+
return c.wrapErr(ErrStopFailed, c.stop(c.cmd))
78+
}
79+
80+
func (c *Cmd) wrapErr(wErr, err error) error {
81+
if err == nil {
82+
return nil
83+
}
84+
return fmt.Errorf("cmd '%s' %w: %w", c.cmd.String(), wErr, err)
85+
}
86+
87+
// WithCommand creates a new command with the given name and arguments.
88+
func WithCommand(name string, args ...string) Opt {
89+
return func(c *Cmd) error {
90+
c.cmd = exec.Command(name, args...)
91+
return nil
92+
}
93+
}
94+
95+
// WithReadyFn allows user to provide custom ready function.
96+
func WithReadyFn(fn func(*exec.Cmd) error) Opt {
97+
return func(c *Cmd) error {
98+
c.ready = fn
99+
return nil
100+
}
101+
}
102+
103+
// WithStopFn allows user to provide custom stop function.
104+
func WithStopFn(fn func(*exec.Cmd) error) Opt {
105+
return func(c *Cmd) error {
106+
c.stop = fn
107+
return nil
108+
}
109+
}
110+
111+
// WithDir sets environment variables for the command.
112+
// By default, the command inherits the environment of the current process and setting this option will override it.
113+
func WithEnv(env ...string) Opt {
114+
return func(c *Cmd) error {
115+
c.cmd.Env = env
116+
return nil
117+
}
118+
}
119+
120+
// WithDir sets the working directory for the command.
121+
func WithDir(dir string) Opt {
122+
return func(c *Cmd) error {
123+
c.cmd.Dir = dir
124+
return nil
125+
}
126+
}
127+
128+
// WithWaitRegexp sets the ready function so that it waits for the command to output a line that matches the given regular expression.
129+
func WithWaitMatchingLine(exp string) Opt {
130+
return func(c *Cmd) error {
131+
re, err := regexp.Compile(exp)
132+
if err != nil {
133+
return err
134+
}
135+
136+
if c.cmd == nil {
137+
return ErrNilCmdRegexp
138+
}
139+
140+
stdout, err := c.cmd.StdoutPipe()
141+
if err != nil {
142+
return err
143+
}
144+
145+
return WithReadyFn(func(cmd *exec.Cmd) error {
146+
scanner := bufio.NewScanner(stdout)
147+
for scanner.Scan() {
148+
if re.Match(scanner.Bytes()) {
149+
return nil
150+
}
151+
}
152+
return errors.Join(ErrNoMatchingLine, scanner.Err())
153+
})(c)
154+
}
155+
}
156+
157+
// WithReadyTimeout overrides the default 30s timeout for the ready function.
158+
func WithReadyTimeout(d time.Duration) Opt {
159+
return func(c *Cmd) error {
160+
c.readyTimeout = d
161+
return nil
162+
}
163+
}
164+
165+
// WithWaitExit sets the ready and stop functions so that ready waits for the command to exit successfully and stop returns nil immediately.
166+
// This is useful for commands that exit on their own and don't need to be stopped manually.
167+
func WithWaitExit() Opt {
168+
return func(c *Cmd) error {
169+
c.ready = func(cmd *exec.Cmd) error { return cmd.Wait() }
170+
c.stop = func(*exec.Cmd) error { return nil }
171+
return nil
172+
}
173+
}
174+
175+
// WithExecCmd allows user to construct the command with custom exec.Cmd.
176+
func WithExecCmd(cmd *exec.Cmd) Opt {
177+
return func(c *Cmd) error {
178+
c.cmd = cmd
179+
return nil
180+
}
181+
}
182+
183+
// StopWithSignal returns a stop function that sends the given signal to the command and waits for it to exit.
184+
// This can be used with WithStopFn to stop the command with a specific signal.
185+
func StopWithSignal(s os.Signal) func(*exec.Cmd) error {
186+
return func(c *exec.Cmd) error {
187+
if c == nil || c.Process == nil {
188+
return nil
189+
}
190+
var err error
191+
if c.Process != nil && c.ProcessState == nil {
192+
err = c.Process.Signal(s)
193+
}
194+
return errors.Join(err, c.Wait())
195+
}
196+
}
197+
198+
func getFnName(fn any) string {
199+
strs := strings.Split((runtime.FuncForPC(reflect.ValueOf(fn).Pointer()).Name()), ".")
200+
return strs[len(strs)-1]
201+
}

‎dep/cmd/cmd_test.go‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
package cmd_test
2+
3+
import (
4+
"os"
5+
"os/exec"
6+
"path/filepath"
7+
"syscall"
8+
"testing"
9+
10+
"github.com/go-tstr/tstr/dep/cmd"
11+
"github.com/go-tstr/tstr/dep/deptest"
12+
"github.com/stretchr/testify/assert"
13+
"github.com/stretchr/testify/require"
14+
)
15+
16+
func TestCmd(t *testing.T) {
17+
testBin := prepareBin(t)
18+
19+
tests := []struct {
20+
name string
21+
fn func() error
22+
cmd *cmd.Cmd
23+
err error
24+
}{
25+
{
26+
name: "MissingCommand",
27+
cmd: cmd.New(),
28+
err: cmd.ErrMissingCmd,
29+
},
30+
{
31+
name: "CommandNotFound",
32+
cmd: cmd.New(
33+
cmd.WithCommand("non-existing-command"),
34+
),
35+
err: cmd.ErrStartFailed,
36+
},
37+
{
38+
name: "WaitForExitError",
39+
cmd: cmd.New(
40+
cmd.WithCommand("go", "foo"),
41+
cmd.WithWaitExit(),
42+
),
43+
err: cmd.ErrReadyFailed,
44+
},
45+
{
46+
name: "WaitForExit",
47+
cmd: cmd.New(
48+
cmd.WithCommand("go", "version"),
49+
cmd.WithWaitExit(),
50+
),
51+
},
52+
{
53+
name: "DefaultStopFn",
54+
cmd: cmd.New(
55+
cmd.WithCommand(testBin),
56+
cmd.WithWaitMatchingLine("Waiting for signal"),
57+
),
58+
},
59+
{
60+
name: "CustomStopFn",
61+
cmd: cmd.New(
62+
cmd.WithCommand(testBin),
63+
cmd.WithWaitMatchingLine("Waiting for signal"),
64+
cmd.WithStopFn(cmd.StopWithSignal(syscall.SIGTERM)),
65+
),
66+
},
67+
{
68+
name: "NoMatchingLine",
69+
cmd: cmd.New(
70+
cmd.WithCommand("go", "version"),
71+
cmd.WithWaitMatchingLine("not matching line"),
72+
),
73+
err: cmd.ErrNoMatchingLine,
74+
},
75+
{
76+
name: "ReadyFailed",
77+
cmd: cmd.New(
78+
cmd.WithCommand("go", "version"),
79+
cmd.WithWaitMatchingLine("not matching line"),
80+
),
81+
err: cmd.ErrReadyFailed,
82+
},
83+
{
84+
name: "ReadyTimeoutExceeded",
85+
cmd: cmd.New(
86+
cmd.WithCommand("sleep", "100"),
87+
cmd.WithReadyFn(blockForever),
88+
cmd.WithReadyTimeout(1),
89+
),
90+
err: cmd.ErrReadyFailed,
91+
},
92+
{
93+
name: "OptionError",
94+
cmd: cmd.New(
95+
cmd.WithWaitMatchingLine("not matching line"),
96+
),
97+
err: cmd.ErrNilCmdRegexp,
98+
},
99+
{
100+
name: "WithDir",
101+
cmd: cmd.New(
102+
cmd.WithCommand("./"+filepath.Base(testBin)),
103+
cmd.WithDir(filepath.Dir(testBin)),
104+
cmd.WithWaitMatchingLine("Waiting for signal"),
105+
),
106+
},
107+
{
108+
name: "WithEnv",
109+
cmd: cmd.New(
110+
cmd.WithCommand("go", "env", "GOPRIVATE"),
111+
cmd.WithEnv("GOPRIVATE=foo"),
112+
cmd.WithWaitMatchingLine("foo"),
113+
),
114+
},
115+
116+
{
117+
name: "WithExecCmd",
118+
cmd: cmd.New(
119+
cmd.WithExecCmd(func() *exec.Cmd { return exec.Command("go", "version") }()),
120+
cmd.WithWaitExit(),
121+
),
122+
},
123+
}
124+
125+
for _, tt := range tests {
126+
t.Run(tt.name, func(t *testing.T) {
127+
deptest.ErrorIs(t, tt.cmd, tt.fn, tt.err)
128+
})
129+
}
130+
}
131+
132+
func blockForever(*exec.Cmd) error {
133+
select {}
134+
}
135+
136+
func prepareBin(t *testing.T) string {
137+
const code = `
138+
package main
139+
140+
import (
141+
"fmt"
142+
"os"
143+
"os/signal"
144+
"syscall"
145+
)
146+
147+
func main() {
148+
c := make(chan os.Signal, 1)
149+
signal.Notify(c, os.Interrupt, syscall.SIGTERM)
150+
fmt.Println("Waiting for signal")
151+
s := <-c
152+
fmt.Println("Got signal:", s)
153+
}`
154+
155+
dir, err := os.MkdirTemp("", "cmd-test-bin_")
156+
require.NoError(t, err)
157+
t.Cleanup(func() { assert.NoError(t, os.RemoveAll(dir)) })
158+
159+
require.NoError(t, os.WriteFile(dir+"/main.go", []byte(code), 0o644))
160+
buildCmd := exec.Command("go", "build", dir+"/main.go")
161+
buildCmd.Dir = dir
162+
require.NoError(t, buildCmd.Run())
163+
return dir + "/main"
164+
}

0 commit comments

Comments
 (0)