Skip to content

Commit 9d3785d

Browse files
mafredrimafredri
authored
test(cli/cliui): make agent tests more robust (#10415)
Fixes #10408
1 parent 2a6fd90 commit 9d3785d

File tree

1 file changed

+101
-43
lines changed

1 file changed

+101
-43
lines changed

‎cli/cliui/agent_test.go

Lines changed: 101 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,14 @@ import (
55
"bytes"
66
"context"
77
"io"
8+
"os"
89
"strings"
910
"sync/atomic"
1011
"testing"
1112
"time"
1213

1314
"github.com/google/uuid"
15+
"github.com/stretchr/testify/assert"
1416
"github.com/stretchr/testify/require"
1517
"golang.org/x/xerrors"
1618

@@ -25,9 +27,31 @@ import (
2527
func TestAgent(t *testing.T) {
2628
t.Parallel()
2729

30+
waitLines := func(t *testing.T, output <-chan string, lines ...string) error {
31+
t.Helper()
32+
33+
var got []string
34+
outerLoop:
35+
for _, want := range lines {
36+
for {
37+
select {
38+
case line := <-output:
39+
got = append(got, line)
40+
if strings.Contains(line, want) {
41+
continue outerLoop
42+
}
43+
case <-time.After(testutil.WaitShort):
44+
assert.Failf(t, "timed out waiting for line", "want: %q; got: %q", want, got)
45+
return xerrors.Errorf("timed out waiting for line: %q; got: %q", want, got)
46+
}
47+
}
48+
}
49+
return nil
50+
}
51+
2852
for _, tc := range []struct {
2953
name string
30-
iter []func(context.Context, *codersdk.WorkspaceAgent, chan []codersdk.WorkspaceAgentLog) error
54+
iter []func(context.Context, *testing.T, *codersdk.WorkspaceAgent, <-chan string, chan []codersdk.WorkspaceAgentLog) error
3155
logs chan []codersdk.WorkspaceAgentLog
3256
opts cliui.AgentOptions
3357
want []string
@@ -38,12 +62,15 @@ func TestAgent(t *testing.T) {
3862
opts: cliui.AgentOptions{
3963
FetchInterval: time.Millisecond,
4064
},
41-
iter: []func(context.Context, *codersdk.WorkspaceAgent, chan []codersdk.WorkspaceAgentLog) error{
42-
func(_ context.Context, agent *codersdk.WorkspaceAgent, _ chan []codersdk.WorkspaceAgentLog) error {
65+
iter: []func(context.Context, *testing.T, *codersdk.WorkspaceAgent, <-chan string, chan []codersdk.WorkspaceAgentLog) error{
66+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, _ <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
4367
agent.Status = codersdk.WorkspaceAgentConnecting
4468
return nil
4569
},
46-
func(_ context.Context, agent *codersdk.WorkspaceAgent, logs chan []codersdk.WorkspaceAgentLog) error {
70+
func(_ context.Context, t *testing.T, agent *codersdk.WorkspaceAgent, output <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
71+
return waitLines(t, output, "⧗ Waiting for the workspace agent to connect")
72+
},
73+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, _ <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
4774
agent.Status = codersdk.WorkspaceAgentConnected
4875
agent.FirstConnectedAt = ptr.Ref(time.Now())
4976
return nil
@@ -62,12 +89,15 @@ func TestAgent(t *testing.T) {
6289
opts: cliui.AgentOptions{
6390
FetchInterval: time.Millisecond,
6491
},
65-
iter: []func(context.Context, *codersdk.WorkspaceAgent, chan []codersdk.WorkspaceAgentLog) error{
66-
func(_ context.Context, agent *codersdk.WorkspaceAgent, _ chan []codersdk.WorkspaceAgentLog) error {
92+
iter: []func(context.Context, *testing.T, *codersdk.WorkspaceAgent, <-chan string, chan []codersdk.WorkspaceAgentLog) error{
93+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, _ <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
6794
agent.Status = codersdk.WorkspaceAgentConnecting
6895
return nil
6996
},
70-
func(_ context.Context, agent *codersdk.WorkspaceAgent, logs chan []codersdk.WorkspaceAgentLog) error {
97+
func(_ context.Context, t *testing.T, agent *codersdk.WorkspaceAgent, output <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
98+
return waitLines(t, output, "⧗ Waiting for the workspace agent to connect")
99+
},
100+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, _ <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
71101
agent.Status = codersdk.WorkspaceAgentConnected
72102
agent.LifecycleState = codersdk.WorkspaceAgentLifecycleStartTimeout
73103
agent.FirstConnectedAt = ptr.Ref(time.Now())
@@ -87,18 +117,24 @@ func TestAgent(t *testing.T) {
87117
opts: cliui.AgentOptions{
88118
FetchInterval: 1 * time.Millisecond,
89119
},
90-
iter: []func(context.Context, *codersdk.WorkspaceAgent, chan []codersdk.WorkspaceAgentLog) error{
91-
func(_ context.Context, agent *codersdk.WorkspaceAgent, _ chan []codersdk.WorkspaceAgentLog) error {
120+
iter: []func(context.Context, *testing.T, *codersdk.WorkspaceAgent, <-chan string, chan []codersdk.WorkspaceAgentLog) error{
121+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, _ <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
92122
agent.Status = codersdk.WorkspaceAgentConnecting
93123
agent.LifecycleState = codersdk.WorkspaceAgentLifecycleStarting
94124
agent.StartedAt = ptr.Ref(time.Now())
95125
return nil
96126
},
97-
func(_ context.Context, agent *codersdk.WorkspaceAgent, _ chan []codersdk.WorkspaceAgentLog) error {
127+
func(_ context.Context, t *testing.T, agent *codersdk.WorkspaceAgent, output <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
128+
return waitLines(t, output, "⧗ Waiting for the workspace agent to connect")
129+
},
130+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, _ <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
98131
agent.Status = codersdk.WorkspaceAgentTimeout
99132
return nil
100133
},
101-
func(_ context.Context, agent *codersdk.WorkspaceAgent, logs chan []codersdk.WorkspaceAgentLog) error {
134+
func(_ context.Context, t *testing.T, agent *codersdk.WorkspaceAgent, output <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
135+
return waitLines(t, output, "The workspace agent is having trouble connecting, wait for it to connect or restart your workspace.")
136+
},
137+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, _ <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
102138
agent.Status = codersdk.WorkspaceAgentConnected
103139
agent.FirstConnectedAt = ptr.Ref(time.Now())
104140
agent.LifecycleState = codersdk.WorkspaceAgentLifecycleReady
@@ -120,8 +156,8 @@ func TestAgent(t *testing.T) {
120156
opts: cliui.AgentOptions{
121157
FetchInterval: 1 * time.Millisecond,
122158
},
123-
iter: []func(context.Context, *codersdk.WorkspaceAgent, chan []codersdk.WorkspaceAgentLog) error{
124-
func(_ context.Context, agent *codersdk.WorkspaceAgent, _ chan []codersdk.WorkspaceAgentLog) error {
159+
iter: []func(context.Context, *testing.T, *codersdk.WorkspaceAgent, <-chan string, chan []codersdk.WorkspaceAgentLog) error{
160+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, _ <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
125161
agent.Status = codersdk.WorkspaceAgentDisconnected
126162
agent.FirstConnectedAt = ptr.Ref(time.Now().Add(-1 * time.Minute))
127163
agent.LastConnectedAt = ptr.Ref(time.Now().Add(-1 * time.Minute))
@@ -131,7 +167,10 @@ func TestAgent(t *testing.T) {
131167
agent.ReadyAt = ptr.Ref(time.Now())
132168
return nil
133169
},
134-
func(_ context.Context, agent *codersdk.WorkspaceAgent, _ chan []codersdk.WorkspaceAgentLog) error {
170+
func(_ context.Context, t *testing.T, agent *codersdk.WorkspaceAgent, output <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
171+
return waitLines(t, output, "⧗ The workspace agent lost connection")
172+
},
173+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, _ <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
135174
agent.Status = codersdk.WorkspaceAgentConnected
136175
agent.DisconnectedAt = nil
137176
agent.LastConnectedAt = ptr.Ref(time.Now())
@@ -151,8 +190,8 @@ func TestAgent(t *testing.T) {
151190
FetchInterval: time.Millisecond,
152191
Wait: true,
153192
},
154-
iter: []func(context.Context, *codersdk.WorkspaceAgent, chan []codersdk.WorkspaceAgentLog) error{
155-
func(_ context.Context, agent *codersdk.WorkspaceAgent, logs chan []codersdk.WorkspaceAgentLog) error {
193+
iter: []func(context.Context, *testing.T, *codersdk.WorkspaceAgent, <-chan string, chan []codersdk.WorkspaceAgentLog) error{
194+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, _ <-chan string, logs chan []codersdk.WorkspaceAgentLog) error {
156195
agent.Status = codersdk.WorkspaceAgentConnected
157196
agent.FirstConnectedAt = ptr.Ref(time.Now())
158197
agent.LifecycleState = codersdk.WorkspaceAgentLifecycleStarting
@@ -170,7 +209,7 @@ func TestAgent(t *testing.T) {
170209
}
171210
return nil
172211
},
173-
func(_ context.Context, agent *codersdk.WorkspaceAgent, logs chan []codersdk.WorkspaceAgentLog) error {
212+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, _ <-chan string, logs chan []codersdk.WorkspaceAgentLog) error {
174213
agent.LifecycleState = codersdk.WorkspaceAgentLifecycleReady
175214
agent.ReadyAt = ptr.Ref(time.Now())
176215
logs <- []codersdk.WorkspaceAgentLog{
@@ -195,8 +234,8 @@ func TestAgent(t *testing.T) {
195234
FetchInterval: time.Millisecond,
196235
Wait: true,
197236
},
198-
iter: []func(context.Context, *codersdk.WorkspaceAgent, chan []codersdk.WorkspaceAgentLog) error{
199-
func(_ context.Context, agent *codersdk.WorkspaceAgent, logs chan []codersdk.WorkspaceAgentLog) error {
237+
iter: []func(context.Context, *testing.T, *codersdk.WorkspaceAgent, <-chan string, chan []codersdk.WorkspaceAgentLog) error{
238+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, output <-chan string, logs chan []codersdk.WorkspaceAgentLog) error {
200239
agent.Status = codersdk.WorkspaceAgentConnected
201240
agent.FirstConnectedAt = ptr.Ref(time.Now())
202241
agent.StartedAt = ptr.Ref(time.Now())
@@ -224,8 +263,8 @@ func TestAgent(t *testing.T) {
224263
opts: cliui.AgentOptions{
225264
FetchInterval: time.Millisecond,
226265
},
227-
iter: []func(context.Context, *codersdk.WorkspaceAgent, chan []codersdk.WorkspaceAgentLog) error{
228-
func(_ context.Context, agent *codersdk.WorkspaceAgent, logs chan []codersdk.WorkspaceAgentLog) error {
266+
iter: []func(context.Context, *testing.T, *codersdk.WorkspaceAgent, <-chan string, chan []codersdk.WorkspaceAgentLog) error{
267+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, output <-chan string, logs chan []codersdk.WorkspaceAgentLog) error {
229268
agent.Status = codersdk.WorkspaceAgentDisconnected
230269
agent.LifecycleState = codersdk.WorkspaceAgentLifecycleOff
231270
return nil
@@ -239,8 +278,8 @@ func TestAgent(t *testing.T) {
239278
FetchInterval: time.Millisecond,
240279
Wait: true,
241280
},
242-
iter: []func(context.Context, *codersdk.WorkspaceAgent, chan []codersdk.WorkspaceAgentLog) error{
243-
func(_ context.Context, agent *codersdk.WorkspaceAgent, logs chan []codersdk.WorkspaceAgentLog) error {
281+
iter: []func(context.Context, *testing.T, *codersdk.WorkspaceAgent, <-chan string, chan []codersdk.WorkspaceAgentLog) error{
282+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, output <-chan string, logs chan []codersdk.WorkspaceAgentLog) error {
244283
agent.Status = codersdk.WorkspaceAgentConnected
245284
agent.FirstConnectedAt = ptr.Ref(time.Now())
246285
agent.LifecycleState = codersdk.WorkspaceAgentLifecycleStarting
@@ -253,7 +292,10 @@ func TestAgent(t *testing.T) {
253292
}
254293
return nil
255294
},
256-
func(_ context.Context, agent *codersdk.WorkspaceAgent, logs chan []codersdk.WorkspaceAgentLog) error {
295+
func(_ context.Context, t *testing.T, agent *codersdk.WorkspaceAgent, output <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
296+
return waitLines(t, output, "Hello world")
297+
},
298+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, _ <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
257299
agent.ReadyAt = ptr.Ref(time.Now())
258300
agent.LifecycleState = codersdk.WorkspaceAgentLifecycleShuttingDown
259301
return nil
@@ -272,12 +314,15 @@ func TestAgent(t *testing.T) {
272314
FetchInterval: time.Millisecond,
273315
Wait: true,
274316
},
275-
iter: []func(context.Context, *codersdk.WorkspaceAgent, chan []codersdk.WorkspaceAgentLog) error{
276-
func(_ context.Context, agent *codersdk.WorkspaceAgent, _ chan []codersdk.WorkspaceAgentLog) error {
317+
iter: []func(context.Context, *testing.T, *codersdk.WorkspaceAgent, <-chan string, chan []codersdk.WorkspaceAgentLog) error{
318+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, _ <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
277319
agent.Status = codersdk.WorkspaceAgentConnecting
278320
return nil
279321
},
280-
func(_ context.Context, agent *codersdk.WorkspaceAgent, _ chan []codersdk.WorkspaceAgentLog) error {
322+
func(_ context.Context, t *testing.T, agent *codersdk.WorkspaceAgent, output <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
323+
return waitLines(t, output, "⧗ Waiting for the workspace agent to connect")
324+
},
325+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, _ <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
281326
return xerrors.New("bad")
282327
},
283328
},
@@ -292,13 +337,16 @@ func TestAgent(t *testing.T) {
292337
FetchInterval: time.Millisecond,
293338
Wait: true,
294339
},
295-
iter: []func(context.Context, *codersdk.WorkspaceAgent, chan []codersdk.WorkspaceAgentLog) error{
296-
func(_ context.Context, agent *codersdk.WorkspaceAgent, _ chan []codersdk.WorkspaceAgentLog) error {
340+
iter: []func(context.Context, *testing.T, *codersdk.WorkspaceAgent, <-chan string, chan []codersdk.WorkspaceAgentLog) error{
341+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, _ <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
297342
agent.Status = codersdk.WorkspaceAgentTimeout
298343
agent.TroubleshootingURL = "https://troubleshoot"
299344
return nil
300345
},
301-
func(_ context.Context, agent *codersdk.WorkspaceAgent, _ chan []codersdk.WorkspaceAgentLog) error {
346+
func(_ context.Context, t *testing.T, agent *codersdk.WorkspaceAgent, output <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
347+
return waitLines(t, output, "The workspace agent is having trouble connecting, wait for it to connect or restart your workspace.")
348+
},
349+
func(_ context.Context, _ *testing.T, agent *codersdk.WorkspaceAgent, output <-chan string, _ chan []codersdk.WorkspaceAgentLog) error {
302350
return xerrors.New("bad")
303351
},
304352
},
@@ -317,21 +365,27 @@ func TestAgent(t *testing.T) {
317365
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
318366
defer cancel()
319367

320-
var buf bytes.Buffer
368+
r, w, err := os.Pipe()
369+
require.NoError(t, err, "create pipe failed")
370+
defer r.Close()
371+
defer w.Close()
372+
321373
agent := codersdk.WorkspaceAgent{
322374
ID: uuid.New(),
323375
Status: codersdk.WorkspaceAgentConnecting,
324376
CreatedAt: time.Now(),
325377
LifecycleState: codersdk.WorkspaceAgentLifecycleCreated,
326378
}
379+
output := make(chan string, 100) // Buffered to avoid blocking, overflow is discarded.
327380
logs := make(chan []codersdk.WorkspaceAgentLog, 1)
328381

329382
cmd := &clibase.Cmd{
330383
Handler: func(inv *clibase.Invocation) error {
331384
tc.opts.Fetch = func(_ context.Context, _ uuid.UUID) (codersdk.WorkspaceAgent, error) {
385+
t.Log("iter", len(tc.iter))
332386
var err error
333387
if len(tc.iter) > 0 {
334-
err = tc.iter[0](ctx, &agent, logs)
388+
err = tc.iter[0](ctx, t, &agent, output, logs)
335389
tc.iter = tc.iter[1:]
336390
}
337391
return agent, err
@@ -352,27 +406,25 @@ func TestAgent(t *testing.T) {
352406
close(fetchLogs)
353407
return fetchLogs, closeFunc(func() error { return nil }), nil
354408
}
355-
err := cliui.Agent(inv.Context(), &buf, uuid.Nil, tc.opts)
409+
err := cliui.Agent(inv.Context(), w, uuid.Nil, tc.opts)
410+
_ = w.Close()
356411
return err
357412
},
358413
}
359414
inv := cmd.Invoke()
360415

361-
w := clitest.StartWithWaiter(t, inv)
362-
if tc.wantErr {
363-
w.RequireError()
364-
} else {
365-
w.RequireSuccess()
366-
}
416+
waiter := clitest.StartWithWaiter(t, inv)
367417

368-
s := bufio.NewScanner(&buf)
418+
s := bufio.NewScanner(r)
369419
for s.Scan() {
370420
line := s.Text()
371421
t.Log(line)
422+
select {
423+
case output <- line:
424+
default:
425+
t.Logf("output overflow: %s", line)
426+
}
372427
if len(tc.want) == 0 {
373-
for i := 0; i < 5; i++ {
374-
t.Log(line)
375-
}
376428
require.Fail(t, "unexpected line", line)
377429
}
378430
require.Contains(t, line, tc.want[0])
@@ -382,6 +434,12 @@ func TestAgent(t *testing.T) {
382434
if len(tc.want) > 0 {
383435
require.Fail(t, "missing lines: "+strings.Join(tc.want, ", "))
384436
}
437+
438+
if tc.wantErr {
439+
waiter.RequireError()
440+
} else {
441+
waiter.RequireSuccess()
442+
}
385443
})
386444
}
387445

0 commit comments

Comments
 (0)