Skip to content

chore: update testutil chan helpers #17408

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Apr 16, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 8 additions & 8 deletions agent/agent_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -110,7 +110,7 @@ func TestAgent_ImmediateClose(t *testing.T) {
})

// wait until the agent has connected and is starting to find races in the startup code
_ = testutil.RequireRecvCtx(ctx, t, client.GetStartup())
_ = testutil.TryReceive(ctx, t, client.GetStartup())
t.Log("Closing Agent")
err := agentUnderTest.Close()
require.NoError(t, err)
Expand Down Expand Up @@ -1700,7 +1700,7 @@ func TestAgent_Lifecycle(t *testing.T) {
// In order to avoid shutting down the agent before it is fully started and triggering
// errors, we'll wait until the agent is fully up. It's a bit hokey, but among the last things the agent starts
// is the stats reporting, so getting a stats report is a good indication the agent is fully up.
_ = testutil.RequireRecvCtx(ctx, t, statsCh)
_ = testutil.TryReceive(ctx, t, statsCh)

err := agent.Close()
require.NoError(t, err, "agent should be closed successfully")
Expand Down Expand Up @@ -1730,7 +1730,7 @@ func TestAgent_Startup(t *testing.T) {
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
Directory: "",
}, 0)
startup := testutil.RequireRecvCtx(ctx, t, client.GetStartup())
startup := testutil.TryReceive(ctx, t, client.GetStartup())
require.Equal(t, "", startup.GetExpandedDirectory())
})

Expand All @@ -1741,7 +1741,7 @@ func TestAgent_Startup(t *testing.T) {
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
Directory: "~",
}, 0)
startup := testutil.RequireRecvCtx(ctx, t, client.GetStartup())
startup := testutil.TryReceive(ctx, t, client.GetStartup())
homeDir, err := os.UserHomeDir()
require.NoError(t, err)
require.Equal(t, homeDir, startup.GetExpandedDirectory())
Expand All @@ -1754,7 +1754,7 @@ func TestAgent_Startup(t *testing.T) {
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
Directory: "coder/coder",
}, 0)
startup := testutil.RequireRecvCtx(ctx, t, client.GetStartup())
startup := testutil.TryReceive(ctx, t, client.GetStartup())
homeDir, err := os.UserHomeDir()
require.NoError(t, err)
require.Equal(t, filepath.Join(homeDir, "coder/coder"), startup.GetExpandedDirectory())
Expand All @@ -1767,7 +1767,7 @@ func TestAgent_Startup(t *testing.T) {
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
Directory: "$HOME",
}, 0)
startup := testutil.RequireRecvCtx(ctx, t, client.GetStartup())
startup := testutil.TryReceive(ctx, t, client.GetStartup())
homeDir, err := os.UserHomeDir()
require.NoError(t, err)
require.Equal(t, homeDir, startup.GetExpandedDirectory())
Expand Down Expand Up @@ -2632,7 +2632,7 @@ done

n := 1
for n <= 5 {
logs := testutil.RequireRecvCtx(ctx, t, logsCh)
logs := testutil.TryReceive(ctx, t, logsCh)
require.NotNil(t, logs)
for _, l := range logs.GetLogs() {
require.Equal(t, fmt.Sprintf("start %d", n), l.GetOutput())
Expand All @@ -2645,7 +2645,7 @@ done

n = 1
for n <= 3000 {
logs := testutil.RequireRecvCtx(ctx, t, logsCh)
logs := testutil.TryReceive(ctx, t, logsCh)
require.NotNil(t, logs)
for _, l := range logs.GetLogs() {
require.Equal(t, fmt.Sprintf("stop %d", n), l.GetOutput())
Expand Down
4 changes: 2 additions & 2 deletions agent/agentscripts/agentscripts_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ func TestExecuteBasic(t *testing.T) {
}}, aAPI.ScriptCompleted)
require.NoError(t, err)
require.NoError(t, runner.Execute(context.Background(), agentscripts.ExecuteAllScripts))
log := testutil.RequireRecvCtx(ctx, t, fLogger.logs)
log := testutil.TryReceive(ctx, t, fLogger.logs)
require.Equal(t, "hello", log.Output)
}

Expand Down Expand Up @@ -136,7 +136,7 @@ func TestScriptReportsTiming(t *testing.T) {
require.NoError(t, runner.Execute(ctx, agentscripts.ExecuteAllScripts))
runner.Close()

log := testutil.RequireRecvCtx(ctx, t, fLogger.logs)
log := testutil.TryReceive(ctx, t, fLogger.logs)
require.Equal(t, "hello", log.Output)

timings := aAPI.GetTimings()
Expand Down
8 changes: 4 additions & 4 deletions agent/apphealth_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ func TestAppHealth_Healthy(t *testing.T) {
mClock.Advance(999 * time.Millisecond).MustWait(ctx) // app2 is now healthy

mClock.Advance(time.Millisecond).MustWait(ctx) // report gets triggered
update := testutil.RequireRecvCtx(ctx, t, fakeAPI.AppHealthCh())
update := testutil.TryReceive(ctx, t, fakeAPI.AppHealthCh())
require.Len(t, update.GetUpdates(), 2)
applyUpdate(t, apps, update)
require.Equal(t, codersdk.WorkspaceAppHealthHealthy, apps[1].Health)
Expand All @@ -101,7 +101,7 @@ func TestAppHealth_Healthy(t *testing.T) {
mClock.Advance(999 * time.Millisecond).MustWait(ctx) // app3 is now healthy

mClock.Advance(time.Millisecond).MustWait(ctx) // report gets triggered
update = testutil.RequireRecvCtx(ctx, t, fakeAPI.AppHealthCh())
update = testutil.TryReceive(ctx, t, fakeAPI.AppHealthCh())
require.Len(t, update.GetUpdates(), 2)
applyUpdate(t, apps, update)
require.Equal(t, codersdk.WorkspaceAppHealthHealthy, apps[1].Health)
Expand Down Expand Up @@ -155,7 +155,7 @@ func TestAppHealth_500(t *testing.T) {
mClock.Advance(999 * time.Millisecond).MustWait(ctx) // 2nd check, crosses threshold
mClock.Advance(time.Millisecond).MustWait(ctx) // 2nd report, sends update

update := testutil.RequireRecvCtx(ctx, t, fakeAPI.AppHealthCh())
update := testutil.TryReceive(ctx, t, fakeAPI.AppHealthCh())
require.Len(t, update.GetUpdates(), 1)
applyUpdate(t, apps, update)
require.Equal(t, codersdk.WorkspaceAppHealthUnhealthy, apps[0].Health)
Expand Down Expand Up @@ -223,7 +223,7 @@ func TestAppHealth_Timeout(t *testing.T) {
timeoutTrap.MustWait(ctx).Release()
mClock.Set(ms(3001)).MustWait(ctx) // report tick, sends changes

update := testutil.RequireRecvCtx(ctx, t, fakeAPI.AppHealthCh())
update := testutil.TryReceive(ctx, t, fakeAPI.AppHealthCh())
require.Len(t, update.GetUpdates(), 1)
applyUpdate(t, apps, update)
require.Equal(t, codersdk.WorkspaceAppHealthUnhealthy, apps[0].Health)
Expand Down
2 changes: 1 addition & 1 deletion agent/checkpoint_internal_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,6 @@ func TestCheckpoint_WaitComplete(t *testing.T) {
errCh <- uut.wait(ctx)
}()
uut.complete(err)
got := testutil.RequireRecvCtx(ctx, t, errCh)
got := testutil.TryReceive(ctx, t, errCh)
require.Equal(t, err, got)
}
26 changes: 13 additions & 13 deletions agent/stats_internal_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -34,14 +34,14 @@ func TestStatsReporter(t *testing.T) {
}()

// initial request to get duration
req := testutil.RequireRecvCtx(ctx, t, fDest.reqs)
req := testutil.TryReceive(ctx, t, fDest.reqs)
require.NotNil(t, req)
require.Nil(t, req.Stats)
interval := time.Second * 34
testutil.RequireSendCtx(ctx, t, fDest.resps, &proto.UpdateStatsResponse{ReportInterval: durationpb.New(interval)})
testutil.RequireSend(ctx, t, fDest.resps, &proto.UpdateStatsResponse{ReportInterval: durationpb.New(interval)})

// call to source to set the callback and interval
gotInterval := testutil.RequireRecvCtx(ctx, t, fSource.period)
gotInterval := testutil.TryReceive(ctx, t, fSource.period)
require.Equal(t, interval, gotInterval)

// callback returning netstats
Expand All @@ -60,7 +60,7 @@ func TestStatsReporter(t *testing.T) {
fSource.callback(time.Now(), time.Now(), netStats, nil)

// collector called to complete the stats
gotNetStats := testutil.RequireRecvCtx(ctx, t, fCollector.calls)
gotNetStats := testutil.TryReceive(ctx, t, fCollector.calls)
require.Equal(t, netStats, gotNetStats)

// while we are collecting the stats, send in two new netStats to simulate
Expand Down Expand Up @@ -94,13 +94,13 @@ func TestStatsReporter(t *testing.T) {

// complete first collection
stats := &proto.Stats{SessionCountJetbrains: 55}
testutil.RequireSendCtx(ctx, t, fCollector.stats, stats)
testutil.RequireSend(ctx, t, fCollector.stats, stats)

// destination called to report the first stats
update := testutil.RequireRecvCtx(ctx, t, fDest.reqs)
update := testutil.TryReceive(ctx, t, fDest.reqs)
require.NotNil(t, update)
require.Equal(t, stats, update.Stats)
testutil.RequireSendCtx(ctx, t, fDest.resps, &proto.UpdateStatsResponse{ReportInterval: durationpb.New(interval)})
testutil.RequireSend(ctx, t, fDest.resps, &proto.UpdateStatsResponse{ReportInterval: durationpb.New(interval)})

// second update -- netStat0 and netStats1 are accumulated and reported
wantNetStats := map[netlogtype.Connection]netlogtype.Counts{
Expand All @@ -115,22 +115,22 @@ func TestStatsReporter(t *testing.T) {
RxBytes: 21,
},
}
gotNetStats = testutil.RequireRecvCtx(ctx, t, fCollector.calls)
gotNetStats = testutil.TryReceive(ctx, t, fCollector.calls)
require.Equal(t, wantNetStats, gotNetStats)
stats = &proto.Stats{SessionCountJetbrains: 66}
testutil.RequireSendCtx(ctx, t, fCollector.stats, stats)
update = testutil.RequireRecvCtx(ctx, t, fDest.reqs)
testutil.RequireSend(ctx, t, fCollector.stats, stats)
update = testutil.TryReceive(ctx, t, fDest.reqs)
require.NotNil(t, update)
require.Equal(t, stats, update.Stats)
interval2 := 27 * time.Second
testutil.RequireSendCtx(ctx, t, fDest.resps, &proto.UpdateStatsResponse{ReportInterval: durationpb.New(interval2)})
testutil.RequireSend(ctx, t, fDest.resps, &proto.UpdateStatsResponse{ReportInterval: durationpb.New(interval2)})

// set the new interval
gotInterval = testutil.RequireRecvCtx(ctx, t, fSource.period)
gotInterval = testutil.TryReceive(ctx, t, fSource.period)
require.Equal(t, interval2, gotInterval)

loopCancel()
err := testutil.RequireRecvCtx(ctx, t, loopErr)
err := testutil.TryReceive(ctx, t, loopErr)
require.NoError(t, err)
}

Expand Down
14 changes: 7 additions & 7 deletions cli/cliui/prompt_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ func TestPrompt(t *testing.T) {
}()
ptty.ExpectMatch("Example")
ptty.WriteLine("hello")
resp := testutil.RequireRecvCtx(ctx, t, msgChan)
resp := testutil.TryReceive(ctx, t, msgChan)
require.Equal(t, "hello", resp)
})

Expand All @@ -54,7 +54,7 @@ func TestPrompt(t *testing.T) {
}()
ptty.ExpectMatch("Example")
ptty.WriteLine("yes")
resp := testutil.RequireRecvCtx(ctx, t, doneChan)
resp := testutil.TryReceive(ctx, t, doneChan)
require.Equal(t, "yes", resp)
})

Expand Down Expand Up @@ -91,7 +91,7 @@ func TestPrompt(t *testing.T) {
doneChan <- resp
}()

resp := testutil.RequireRecvCtx(ctx, t, doneChan)
resp := testutil.TryReceive(ctx, t, doneChan)
require.Equal(t, "yes", resp)
// Close the reader to end the io.Copy
require.NoError(t, ptty.Close(), "close eof reader")
Expand All @@ -115,7 +115,7 @@ func TestPrompt(t *testing.T) {
}()
ptty.ExpectMatch("Example")
ptty.WriteLine("{}")
resp := testutil.RequireRecvCtx(ctx, t, doneChan)
resp := testutil.TryReceive(ctx, t, doneChan)
require.Equal(t, "{}", resp)
})

Expand All @@ -133,7 +133,7 @@ func TestPrompt(t *testing.T) {
}()
ptty.ExpectMatch("Example")
ptty.WriteLine("{a")
resp := testutil.RequireRecvCtx(ctx, t, doneChan)
resp := testutil.TryReceive(ctx, t, doneChan)
require.Equal(t, "{a", resp)
})

Expand All @@ -153,7 +153,7 @@ func TestPrompt(t *testing.T) {
ptty.WriteLine(`{
"test": "wow"
}`)
resp := testutil.RequireRecvCtx(ctx, t, doneChan)
resp := testutil.TryReceive(ctx, t, doneChan)
require.Equal(t, `{"test":"wow"}`, resp)
})

Expand All @@ -178,7 +178,7 @@ func TestPrompt(t *testing.T) {
}()
ptty.ExpectMatch("Example")
ptty.WriteLine("foo\nbar\nbaz\n\n\nvalid\n")
resp := testutil.RequireRecvCtx(ctx, t, doneChan)
resp := testutil.TryReceive(ctx, t, doneChan)
require.Equal(t, "valid", resp)
})
}
Expand Down
16 changes: 8 additions & 8 deletions cli/portforward_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -192,8 +192,8 @@ func TestPortForward(t *testing.T) {
require.ErrorIs(t, err, context.Canceled)

flushCtx := testutil.Context(t, testutil.WaitShort)
testutil.RequireSendCtx(flushCtx, t, wuTick, dbtime.Now())
_ = testutil.RequireRecvCtx(flushCtx, t, wuFlush)
testutil.RequireSend(flushCtx, t, wuTick, dbtime.Now())
_ = testutil.TryReceive(flushCtx, t, wuFlush)
updated, err := client.Workspace(context.Background(), workspace.ID)
require.NoError(t, err)
require.Greater(t, updated.LastUsedAt, workspace.LastUsedAt)
Expand Down Expand Up @@ -247,8 +247,8 @@ func TestPortForward(t *testing.T) {
require.ErrorIs(t, err, context.Canceled)

flushCtx := testutil.Context(t, testutil.WaitShort)
testutil.RequireSendCtx(flushCtx, t, wuTick, dbtime.Now())
_ = testutil.RequireRecvCtx(flushCtx, t, wuFlush)
testutil.RequireSend(flushCtx, t, wuTick, dbtime.Now())
_ = testutil.TryReceive(flushCtx, t, wuFlush)
updated, err := client.Workspace(context.Background(), workspace.ID)
require.NoError(t, err)
require.Greater(t, updated.LastUsedAt, workspace.LastUsedAt)
Expand Down Expand Up @@ -315,8 +315,8 @@ func TestPortForward(t *testing.T) {
require.ErrorIs(t, err, context.Canceled)

flushCtx := testutil.Context(t, testutil.WaitShort)
testutil.RequireSendCtx(flushCtx, t, wuTick, dbtime.Now())
_ = testutil.RequireRecvCtx(flushCtx, t, wuFlush)
testutil.RequireSend(flushCtx, t, wuTick, dbtime.Now())
_ = testutil.TryReceive(flushCtx, t, wuFlush)
updated, err := client.Workspace(context.Background(), workspace.ID)
require.NoError(t, err)
require.Greater(t, updated.LastUsedAt, workspace.LastUsedAt)
Expand Down Expand Up @@ -372,8 +372,8 @@ func TestPortForward(t *testing.T) {
require.ErrorIs(t, err, context.Canceled)

flushCtx := testutil.Context(t, testutil.WaitShort)
testutil.RequireSendCtx(flushCtx, t, wuTick, dbtime.Now())
_ = testutil.RequireRecvCtx(flushCtx, t, wuFlush)
testutil.RequireSend(flushCtx, t, wuTick, dbtime.Now())
_ = testutil.TryReceive(flushCtx, t, wuFlush)
updated, err := client.Workspace(context.Background(), workspace.ID)
require.NoError(t, err)
require.Greater(t, updated.LastUsedAt, workspace.LastUsedAt)
Expand Down
14 changes: 7 additions & 7 deletions cli/ssh_internal_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,7 @@ func TestCloserStack_Empty(t *testing.T) {
defer close(closed)
uut.close(nil)
}()
testutil.RequireRecvCtx(ctx, t, closed)
testutil.TryReceive(ctx, t, closed)
}

func TestCloserStack_Context(t *testing.T) {
Expand Down Expand Up @@ -157,7 +157,7 @@ func TestCloserStack_CloseAfterContext(t *testing.T) {
err := uut.push("async", ac)
require.NoError(t, err)
cancel()
testutil.RequireRecvCtx(testCtx, t, ac.started)
testutil.TryReceive(testCtx, t, ac.started)

closed := make(chan struct{})
go func() {
Expand All @@ -174,7 +174,7 @@ func TestCloserStack_CloseAfterContext(t *testing.T) {
}

ac.complete()
testutil.RequireRecvCtx(testCtx, t, closed)
testutil.TryReceive(testCtx, t, closed)
}

func TestCloserStack_Timeout(t *testing.T) {
Expand Down Expand Up @@ -204,20 +204,20 @@ func TestCloserStack_Timeout(t *testing.T) {
}()
trap.MustWait(ctx).Release()
// top starts right away, but it hangs
testutil.RequireRecvCtx(ctx, t, ac[2].started)
testutil.TryReceive(ctx, t, ac[2].started)
// timer pops and we start the middle one
mClock.Advance(gracefulShutdownTimeout).MustWait(ctx)
testutil.RequireRecvCtx(ctx, t, ac[1].started)
testutil.TryReceive(ctx, t, ac[1].started)

// middle one finishes
ac[1].complete()
// bottom starts, but also hangs
testutil.RequireRecvCtx(ctx, t, ac[0].started)
testutil.TryReceive(ctx, t, ac[0].started)

// timer has to pop twice to time out.
mClock.Advance(gracefulShutdownTimeout).MustWait(ctx)
mClock.Advance(gracefulShutdownTimeout).MustWait(ctx)
testutil.RequireRecvCtx(ctx, t, closed)
testutil.TryReceive(ctx, t, closed)
}

type fakeCloser struct {
Expand Down
10 changes: 5 additions & 5 deletions cli/ssh_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -271,12 +271,12 @@ func TestSSH(t *testing.T) {
}

// Allow one build to complete.
testutil.RequireSendCtx(ctx, t, buildPause, true)
testutil.RequireRecvCtx(ctx, t, buildDone)
testutil.RequireSend(ctx, t, buildPause, true)
testutil.TryReceive(ctx, t, buildDone)

// Allow the remaining builds to continue.
for i := 0; i < len(ptys)-1; i++ {
testutil.RequireSendCtx(ctx, t, buildPause, false)
testutil.RequireSend(ctx, t, buildPause, false)
}

var foundConflict int
Expand Down Expand Up @@ -1017,14 +1017,14 @@ func TestSSH(t *testing.T) {
}
}()

msg := testutil.RequireRecvCtx(ctx, t, msgs)
msg := testutil.TryReceive(ctx, t, msgs)
require.Equal(t, "test", msg)
close(success)
fsn.Notify()
<-cmdDone
fsn.AssertStopped()
// wait for dial goroutine to complete
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)

// wait for the remote socket to get cleaned up before retrying,
// because cleaning up the socket happens asynchronously, and we
Expand Down
Loading
Loading