-
Notifications
You must be signed in to change notification settings - Fork 1.9k
/
cmd.go
81 lines (65 loc) · 1.87 KB
/
cmd.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
package test
import (
"context"
"os"
"testing"
"time"
"github.com/docker/cli/cli/command"
"github.com/docker/cli/cli/streams"
"github.com/spf13/cobra"
"gotest.tools/v3/assert"
)
func TerminatePrompt(ctx context.Context, t *testing.T, cmd *cobra.Command, cli *FakeCli) {
t.Helper()
errChan := make(chan error)
defer close(errChan)
// wrap the out stream to detect when the prompt is ready
writerHookChan := make(chan struct{})
defer close(writerHookChan)
outStream := streams.NewOut(NewWriterWithHook(cli.OutBuffer(), func(p []byte) {
writerHookChan <- struct{}{}
}))
cli.SetOut(outStream)
r, _, err := os.Pipe()
assert.NilError(t, err)
cli.SetIn(streams.NewIn(r))
notifyCtx, notifyCancel := context.WithCancel(ctx)
t.Cleanup(notifyCancel)
go func() {
errChan <- cmd.ExecuteContext(notifyCtx)
}()
writeCtx, writeCancel := context.WithTimeout(ctx, 100*time.Millisecond)
defer writeCancel()
// wait for the prompt to be ready
select {
case <-writeCtx.Done():
t.Fatalf("command %s did not write prompt to stdout", cmd.Name())
case <-writerHookChan:
// drain the channel for future buffer writes
go func() {
for {
select {
case <-ctx.Done():
return
case <-writerHookChan:
}
}
}()
}
assert.Check(t, cli.OutBuffer().Len() > 0)
// a small delay to ensure the plugin is prompting
time.Sleep(100 * time.Microsecond)
errCtx, errCancel := context.WithTimeout(ctx, 100*time.Millisecond)
defer errCancel()
// sigint and sigterm are caught by the prompt
// this allows us to gracefully exit the prompt with a 0 exit code
notifyCancel()
select {
case <-errCtx.Done():
t.Logf("command stdout:\n%s\n", cli.OutBuffer().String())
t.Logf("command stderr:\n%s\n", cli.ErrBuffer().String())
t.Fatalf("command %s did not return after SIGINT", cmd.Name())
case err := <-errChan:
assert.ErrorIs(t, err, command.ErrPromptTerminated)
}
}