shishantbiswas/coder
1
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-06-02 20:48:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore: update testutil chan helpers (#17408)

Browse Source
This commit is contained in:
ケイラ
2025-04-16 09:37:09 -07:00
committed by GitHub
parent 2a76f5028e
commit f670bc31f5
45 changed files with 582 additions and 559 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agent/agent_test.go
+8 -8
View File
@@ -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)
@@ -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")
@@ -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())
})
@@ -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())
@@ -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())
@@ -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())
@@ -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())
@@ -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())
agent/agentscripts/agentscripts_test.go
+2 -2
View File
@@ -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)
}
@@ -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()
agent/apphealth_test.go
+4 -4
View File
@@ -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)
@@ -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)
@@ -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)
@@ -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)
agent/checkpoint_internal_test.go
+1 -1
View File
@@ -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)
}
agent/stats_internal_test.go
+13 -13
View File
@@ -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
@@ -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
@@ -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{
@@ -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)
}
cli/cliui/prompt_test.go
+7 -7
View File
@@ -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)
})
@@ -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)
})
@@ -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")
@@ -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)
})
@@ -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)
})
@@ -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)
})
@@ -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)
})
}
cli/portforward_test.go
+8 -8
View File
@@ -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)
@@ -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)
@@ -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)
@@ -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)
cli/ssh_internal_test.go
+7 -7
View File
@@ -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) {
@@ -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() {
@@ -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) {
@@ -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 {
cli/ssh_test.go
+5 -5
View File
@@ -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
@@ -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
cli/start_test.go
+3 -3
View File
@@ -408,7 +408,7 @@ func TestStart_AlreadyRunning(t *testing.T) {
}()
pty.ExpectMatch("workspace is already running")
_ = testutil.RequireRecvCtx(ctx, t, doneChan)
_ = testutil.TryReceive(ctx, t, doneChan)
}
func TestStart_Starting(t *testing.T) {
@@ -441,7 +441,7 @@ func TestStart_Starting(t *testing.T) {
_ = dbfake.JobComplete(t, store, r.Build.JobID).Pubsub(ps).Do()
pty.ExpectMatch("workspace has been started")
_ = testutil.RequireRecvCtx(ctx, t, doneChan)
_ = testutil.TryReceive(ctx, t, doneChan)
}
func TestStart_NoWait(t *testing.T) {
@@ -474,5 +474,5 @@ func TestStart_NoWait(t *testing.T) {
}()
pty.ExpectMatch("workspace has been started in no-wait mode")
_ = testutil.RequireRecvCtx(ctx, t, doneChan)
_ = testutil.TryReceive(ctx, t, doneChan)
}
cli/update_test.go
+10 -10
View File
@@ -345,7 +345,7 @@ func TestUpdateValidateRichParameters(t *testing.T) {
pty.ExpectMatch("does not match")
pty.ExpectMatch("> Enter a value (default: \"\"): ")
pty.WriteLine("abc")
_ = testutil.RequireRecvCtx(ctx, t, doneChan)
_ = testutil.TryReceive(ctx, t, doneChan)
})
t.Run("ValidateNumber", func(t *testing.T) {
@@ -391,7 +391,7 @@ func TestUpdateValidateRichParameters(t *testing.T) {
pty.ExpectMatch("is not a number")
pty.ExpectMatch("> Enter a value (default: \"\"): ")
pty.WriteLine("8")
_ = testutil.RequireRecvCtx(ctx, t, doneChan)
_ = testutil.TryReceive(ctx, t, doneChan)
})
t.Run("ValidateBool", func(t *testing.T) {
@@ -437,7 +437,7 @@ func TestUpdateValidateRichParameters(t *testing.T) {
pty.ExpectMatch("boolean value can be either \"true\" or \"false\"")
pty.ExpectMatch("> Enter a value (default: \"\"): ")
pty.WriteLine("false")
_ = testutil.RequireRecvCtx(ctx, t, doneChan)
_ = testutil.TryReceive(ctx, t, doneChan)
})
t.Run("RequiredParameterAdded", func(t *testing.T) {
@@ -508,7 +508,7 @@ func TestUpdateValidateRichParameters(t *testing.T) {
pty.WriteLine(value)
}
}
_ = testutil.RequireRecvCtx(ctx, t, doneChan)
_ = testutil.TryReceive(ctx, t, doneChan)
})
t.Run("OptionalParameterAdded", func(t *testing.T) {
@@ -568,7 +568,7 @@ func TestUpdateValidateRichParameters(t *testing.T) {
}()
pty.ExpectMatch("Planning workspace...")
_ = testutil.RequireRecvCtx(ctx, t, doneChan)
_ = testutil.TryReceive(ctx, t, doneChan)
})
t.Run("ParameterOptionChanged", func(t *testing.T) {
@@ -640,7 +640,7 @@ func TestUpdateValidateRichParameters(t *testing.T) {
}
}
_ = testutil.RequireRecvCtx(ctx, t, doneChan)
_ = testutil.TryReceive(ctx, t, doneChan)
})
t.Run("ParameterOptionDisappeared", func(t *testing.T) {
@@ -713,7 +713,7 @@ func TestUpdateValidateRichParameters(t *testing.T) {
}
}
_ = testutil.RequireRecvCtx(ctx, t, doneChan)
_ = testutil.TryReceive(ctx, t, doneChan)
})
t.Run("ParameterOptionFailsMonotonicValidation", func(t *testing.T) {
@@ -770,7 +770,7 @@ func TestUpdateValidateRichParameters(t *testing.T) {
pty.ExpectMatch(match)
}
_ = testutil.RequireRecvCtx(ctx, t, doneChan)
_ = testutil.TryReceive(ctx, t, doneChan)
})
t.Run("ImmutableRequiredParameterExists_MutableRequiredParameterAdded", func(t *testing.T) {
@@ -838,7 +838,7 @@ func TestUpdateValidateRichParameters(t *testing.T) {
}
}
_ = testutil.RequireRecvCtx(ctx, t, doneChan)
_ = testutil.TryReceive(ctx, t, doneChan)
})
t.Run("MutableRequiredParameterExists_ImmutableRequiredParameterAdded", func(t *testing.T) {
@@ -910,6 +910,6 @@ func TestUpdateValidateRichParameters(t *testing.T) {
}
}
_ = testutil.RequireRecvCtx(ctx, t, doneChan)
_ = testutil.TryReceive(ctx, t, doneChan)
})
}
cli/usercreate_test.go
+2 -2
View File
@@ -39,7 +39,7 @@ func TestUserCreate(t *testing.T) {
pty.ExpectMatch(match)
pty.WriteLine(value)
}
_ = testutil.RequireRecvCtx(ctx, t, doneChan)
_ = testutil.TryReceive(ctx, t, doneChan)
created, err := client.User(ctx, matches[1])
require.NoError(t, err)
assert.Equal(t, matches[1], created.Username)
@@ -72,7 +72,7 @@ func TestUserCreate(t *testing.T) {
pty.ExpectMatch(match)
pty.WriteLine(value)
}
_ = testutil.RequireRecvCtx(ctx, t, doneChan)
_ = testutil.TryReceive(ctx, t, doneChan)
created, err := client.User(ctx, matches[1])
require.NoError(t, err)
assert.Equal(t, matches[1], created.Username)
coderd/autobuild/lifecycle_executor_test.go
+2 -2
View File
@@ -400,7 +400,7 @@ func TestExecutorAutostartUserSuspended(t *testing.T) {
}()
// Then: nothing should happen
stats := testutil.RequireRecvCtx(ctx, t, statsCh)
stats := testutil.TryReceive(ctx, t, statsCh)
assert.Len(t, stats.Errors, 0)
assert.Len(t, stats.Transitions, 0)
}
@@ -1167,7 +1167,7 @@ func TestNotifications(t *testing.T) {
// Wait for workspace to become dormant
notifyEnq.Clear()
ticker <- workspace.LastUsedAt.Add(timeTilDormant * 3)
_ = testutil.RequireRecvCtx(testutil.Context(t, testutil.WaitShort), t, statCh)
_ = testutil.TryReceive(testutil.Context(t, testutil.WaitShort), t, statCh)
// Check that the workspace is dormant
workspace = coderdtest.MustWorkspace(t, client, workspace.ID)
coderd/database/pubsub/pubsub_internal_test.go
+4 -4
View File
@@ -160,19 +160,19 @@ func TestPubSub_DoesntBlockNotify(t *testing.T) {
assert.NoError(t, err)
cancels <- subCancel
}()
subCancel := testutil.RequireRecvCtx(ctx, t, cancels)
subCancel := testutil.TryReceive(ctx, t, cancels)
cancelDone := make(chan struct{})
go func() {
defer close(cancelDone)
subCancel()
}()
testutil.RequireRecvCtx(ctx, t, cancelDone)
testutil.TryReceive(ctx, t, cancelDone)
closeErrs := make(chan error)
go func() {
closeErrs <- uut.Close()
}()
err := testutil.RequireRecvCtx(ctx, t, closeErrs)
err := testutil.TryReceive(ctx, t, closeErrs)
require.NoError(t, err)
}
@@ -221,7 +221,7 @@ func TestPubSub_DoesntRaceListenUnlisten(t *testing.T) {
}
close(start)
for range numEvents * 2 {
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
for i := range events {
fListener.requireIsListening(t, events[i])
coderd/database/pubsub/pubsub_test.go
+7 -7
View File
@@ -60,7 +60,7 @@ func TestPGPubsub_Metrics(t *testing.T) {
err := uut.Publish(event, []byte(data))
assert.NoError(t, err)
}()
_ = testutil.RequireRecvCtx(ctx, t, messageChannel)
_ = testutil.TryReceive(ctx, t, messageChannel)
require.Eventually(t, func() bool {
latencyBytes := gatherCount * pubsub.LatencyMessageLength
@@ -96,8 +96,8 @@ func TestPGPubsub_Metrics(t *testing.T) {
assert.NoError(t, err)
}()
// should get 2 messages because we have 2 subs
_ = testutil.RequireRecvCtx(ctx, t, messageChannel)
_ = testutil.RequireRecvCtx(ctx, t, messageChannel)
_ = testutil.TryReceive(ctx, t, messageChannel)
_ = testutil.TryReceive(ctx, t, messageChannel)
require.Eventually(t, func() bool {
latencyBytes := gatherCount * pubsub.LatencyMessageLength
@@ -167,10 +167,10 @@ func TestPGPubsubDriver(t *testing.T) {
require.NoError(t, err)
// wait for the message
_ = testutil.RequireRecvCtx(ctx, t, gotChan)
_ = testutil.TryReceive(ctx, t, gotChan)
// read out first connection
firstConn := testutil.RequireRecvCtx(ctx, t, subDriver.Connections)
firstConn := testutil.TryReceive(ctx, t, subDriver.Connections)
// drop the underlying connection being used by the pubsub
// the pq.Listener should reconnect and repopulate it's listeners
@@ -179,7 +179,7 @@ func TestPGPubsubDriver(t *testing.T) {
require.NoError(t, err)
// wait for the reconnect
_ = testutil.RequireRecvCtx(ctx, t, subDriver.Connections)
_ = testutil.TryReceive(ctx, t, subDriver.Connections)
// we need to sleep because the raw connection notification
// is sent before the pq.Listener can reestablish it's listeners
time.Sleep(1 * time.Second)
@@ -189,5 +189,5 @@ func TestPGPubsubDriver(t *testing.T) {
require.NoError(t, err)
// wait for the message on the old subscription
_ = testutil.RequireRecvCtx(ctx, t, gotChan)
_ = testutil.TryReceive(ctx, t, gotChan)
}
coderd/database/pubsub/watchdog_test.go
+7 -7
View File
@@ -37,7 +37,7 @@ func TestWatchdog_NoTimeout(t *testing.T) {
// we subscribe after starting the timer, so we know the timer also starts
// from the baseline.
sub := testutil.RequireRecvCtx(ctx, t, fPS.subs)
sub := testutil.TryReceive(ctx, t, fPS.subs)
require.Equal(t, pubsub.EventPubsubWatchdog, sub.event)
// 5 min / 15 sec = 20, so do 21 ticks
@@ -45,7 +45,7 @@ func TestWatchdog_NoTimeout(t *testing.T) {
d, w := mClock.AdvanceNext()
w.MustWait(ctx)
require.LessOrEqual(t, d, 15*time.Second)
p := testutil.RequireRecvCtx(ctx, t, fPS.pubs)
p := testutil.TryReceive(ctx, t, fPS.pubs)
require.Equal(t, pubsub.EventPubsubWatchdog, p)
mClock.Advance(30 * time.Millisecond). // reasonable round-trip
MustWait(ctx)
@@ -67,7 +67,7 @@ func TestWatchdog_NoTimeout(t *testing.T) {
sc, err := subTrap.Wait(ctx) // timer.Stop() called
require.NoError(t, err)
sc.Release()
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
require.NoError(t, err)
}
@@ -93,7 +93,7 @@ func TestWatchdog_Timeout(t *testing.T) {
// we subscribe after starting the timer, so we know the timer also starts
// from the baseline.
sub := testutil.RequireRecvCtx(ctx, t, fPS.subs)
sub := testutil.TryReceive(ctx, t, fPS.subs)
require.Equal(t, pubsub.EventPubsubWatchdog, sub.event)
// 5 min / 15 sec = 20, so do 19 ticks without timing out
@@ -101,7 +101,7 @@ func TestWatchdog_Timeout(t *testing.T) {
d, w := mClock.AdvanceNext()
w.MustWait(ctx)
require.LessOrEqual(t, d, 15*time.Second)
p := testutil.RequireRecvCtx(ctx, t, fPS.pubs)
p := testutil.TryReceive(ctx, t, fPS.pubs)
require.Equal(t, pubsub.EventPubsubWatchdog, p)
mClock.Advance(30 * time.Millisecond). // reasonable round-trip
MustWait(ctx)
@@ -117,9 +117,9 @@ func TestWatchdog_Timeout(t *testing.T) {
d, w := mClock.AdvanceNext()
w.MustWait(ctx)
require.LessOrEqual(t, d, 15*time.Second)
p := testutil.RequireRecvCtx(ctx, t, fPS.pubs)
p := testutil.TryReceive(ctx, t, fPS.pubs)
require.Equal(t, pubsub.EventPubsubWatchdog, p)
testutil.RequireRecvCtx(ctx, t, uut.Timeout())
testutil.TryReceive(ctx, t, uut.Timeout())
err = uut.Close()
require.NoError(t, err)
coderd/entitlements/entitlements_test.go
+3 -3
View File
@@ -78,7 +78,7 @@ func TestUpdate(t *testing.T) {
})
errCh <- err
}()
testutil.RequireRecvCtx(ctx, t, fetchStarted)
testutil.TryReceive(ctx, t, fetchStarted)
require.False(t, set.Enabled(codersdk.FeatureMultipleOrganizations))
// start a second update while the first one is in progress
go func() {
@@ -97,9 +97,9 @@ func TestUpdate(t *testing.T) {
errCh <- err
}()
close(firstDone)
err := testutil.RequireRecvCtx(ctx, t, errCh)
err := testutil.TryReceive(ctx, t, errCh)
require.NoError(t, err)
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
require.NoError(t, err)
require.True(t, set.Enabled(codersdk.FeatureMultipleOrganizations))
require.True(t, set.Enabled(codersdk.FeatureAppearance))
coderd/httpmw/loggermw/logger_internal_test.go
+2 -2
View File
@@ -146,7 +146,7 @@ func TestLoggerMiddleware_WebSocket(t *testing.T) {
defer conn.Close(websocket.StatusNormalClosure, "")
// Wait for the log from within the handler
newEntry := testutil.RequireRecvCtx(ctx, t, sink.newEntries)
newEntry := testutil.TryReceive(ctx, t, sink.newEntries)
require.Equal(t, newEntry.Message, "GET")
// Signal the websocket handler to return (and read to handle the close frame)
@@ -155,7 +155,7 @@ func TestLoggerMiddleware_WebSocket(t *testing.T) {
require.ErrorAs(t, err, &websocket.CloseError{}, "websocket read should fail with close error")
// Wait for the request to finish completely and verify we only logged once
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
require.Len(t, sink.entries, 1, "log was written twice")
}
coderd/notifications/manager_test.go
+1 -1
View File
@@ -155,7 +155,7 @@ func TestBuildPayload(t *testing.T) {
require.NoError(t, err)
// THEN: expect that a payload will be constructed and have the expected values
payload := testutil.RequireRecvCtx(ctx, t, interceptor.payload)
payload := testutil.TryReceive(ctx, t, interceptor.payload)
require.Len(t, payload.Actions, 1)
require.Equal(t, label, payload.Actions[0].Label)
require.Equal(t, url, payload.Actions[0].URL)
coderd/notifications/metrics_test.go
+2 -2
View File
@@ -300,9 +300,9 @@ func TestPendingUpdatesMetric(t *testing.T) {
mClock.Advance(cfg.StoreSyncInterval.Value() - cfg.FetchInterval.Value()).MustWait(ctx)
// Wait until we intercept the calls to sync the pending updates to the store.
success := testutil.RequireRecvCtx(testutil.Context(t, testutil.WaitShort), t, interceptor.updateSuccess)
success := testutil.TryReceive(testutil.Context(t, testutil.WaitShort), t, interceptor.updateSuccess)
require.EqualValues(t, 2, success)
failure := testutil.RequireRecvCtx(testutil.Context(t, testutil.WaitShort), t, interceptor.updateFailure)
failure := testutil.TryReceive(testutil.Context(t, testutil.WaitShort), t, interceptor.updateFailure)
require.EqualValues(t, 2, failure)
// Validate that the store synced the expected number of updates.
coderd/notifications/notifications_test.go
+5 -5
View File
@@ -260,7 +260,7 @@ func TestWebhookDispatch(t *testing.T) {
mgr.Run(ctx)
// THEN: the webhook is received by the mock server and has the expected contents
payload := testutil.RequireRecvCtx(testutil.Context(t, testutil.WaitShort), t, sent)
payload := testutil.TryReceive(testutil.Context(t, testutil.WaitShort), t, sent)
require.EqualValues(t, "1.1", payload.Version)
require.Equal(t, msgID[0], payload.MsgID)
require.Equal(t, payload.Payload.Labels, input)
@@ -350,8 +350,8 @@ func TestBackpressure(t *testing.T) {
// one batch of dispatches is sent
for range batchSize {
call := testutil.RequireRecvCtx(ctx, t, handler.calls)
testutil.RequireSendCtx(ctx, t, call.result, dispatchResult{
call := testutil.TryReceive(ctx, t, handler.calls)
testutil.RequireSend(ctx, t, call.result, dispatchResult{
retryable: false,
err: nil,
})
@@ -402,7 +402,7 @@ func TestBackpressure(t *testing.T) {
// The batch completes
w.MustWait(ctx)
require.NoError(t, testutil.RequireRecvCtx(ctx, t, stopErr))
require.NoError(t, testutil.TryReceive(ctx, t, stopErr))
require.EqualValues(t, batchSize, storeInterceptor.sent.Load()+storeInterceptor.failed.Load())
}
@@ -1808,7 +1808,7 @@ func TestCustomNotificationMethod(t *testing.T) {
// THEN: the notification should be received by the custom dispatch method
mgr.Run(ctx)
receivedMsgID := testutil.RequireRecvCtx(ctx, t, received)
receivedMsgID := testutil.TryReceive(ctx, t, received)
require.Equal(t, msgID[0].String(), receivedMsgID.String())
// Ensure no messages received by default method (SMTP):
coderd/provisionerdserver/provisionerdserver_test.go
+1 -1
View File
@@ -118,7 +118,7 @@ func TestHeartbeat(t *testing.T) {
})
for i := 0; i < numBeats; i++ {
testutil.RequireRecvCtx(ctx, t, heartbeatChan)
testutil.TryReceive(ctx, t, heartbeatChan)
}
// goleak.VerifyTestMain ensures that the heartbeat goroutine does not leak
}
coderd/rbac/authz_test.go
+4 -4
View File
@@ -362,7 +362,7 @@ func TestCache(t *testing.T) {
authOut = make(chan error, 1) // buffered to not block
authorizeFunc = func(ctx context.Context, subject rbac.Subject, action policy.Action, object rbac.Object) error {
// Just return what you're told.
return testutil.RequireRecvCtx(ctx, t, authOut)
return testutil.TryReceive(ctx, t, authOut)
}
ma = &rbac.MockAuthorizer{AuthorizeFunc: authorizeFunc}
rec = &coderdtest.RecordingAuthorizer{Wrapped: ma}
@@ -371,12 +371,12 @@ func TestCache(t *testing.T) {
)
// First call will result in a transient error. This should not be cached.
testutil.RequireSendCtx(ctx, t, authOut, context.Canceled)
testutil.RequireSend(ctx, t, authOut, context.Canceled)
err := authz.Authorize(ctx, subj, action, obj)
assert.ErrorIs(t, err, context.Canceled)
// A subsequent call should still hit the authorizer.
testutil.RequireSendCtx(ctx, t, authOut, nil)
testutil.RequireSend(ctx, t, authOut, nil)
err = authz.Authorize(ctx, subj, action, obj)
assert.NoError(t, err)
// This should be cached and not hit the wrapped authorizer again.
@@ -387,7 +387,7 @@ func TestCache(t *testing.T) {
subj, obj, action = coderdtest.RandomRBACSubject(), coderdtest.RandomRBACObject(), coderdtest.RandomRBACAction()
// A third will be a legit error
testutil.RequireSendCtx(ctx, t, authOut, assert.AnError)
testutil.RequireSend(ctx, t, authOut, assert.AnError)
err = authz.Authorize(ctx, subj, action, obj)
assert.EqualError(t, err, assert.AnError.Error())
// This should be cached and not hit the wrapped authorizer again.
coderd/templateversions_test.go
+3 -3
View File
@@ -2172,7 +2172,7 @@ func TestTemplateVersionDynamicParameters(t *testing.T) {
previews := stream.Chan()
// Should automatically send a form state with all defaulted/empty values
preview := testutil.RequireRecvCtx(ctx, t, previews)
preview := testutil.TryReceive(ctx, t, previews)
require.Empty(t, preview.Diagnostics)
require.Equal(t, "group", preview.Parameters[0].Name)
require.True(t, preview.Parameters[0].Value.Valid())
@@ -2184,7 +2184,7 @@ func TestTemplateVersionDynamicParameters(t *testing.T) {
Inputs: map[string]string{"group": "Bloob"},
})
require.NoError(t, err)
preview = testutil.RequireRecvCtx(ctx, t, previews)
preview = testutil.TryReceive(ctx, t, previews)
require.Equal(t, 1, preview.ID)
require.Empty(t, preview.Diagnostics)
require.Equal(t, "group", preview.Parameters[0].Name)
@@ -2197,7 +2197,7 @@ func TestTemplateVersionDynamicParameters(t *testing.T) {
Inputs: map[string]string{},
})
require.NoError(t, err)
preview = testutil.RequireRecvCtx(ctx, t, previews)
preview = testutil.TryReceive(ctx, t, previews)
require.Equal(t, 3, preview.ID)
require.Empty(t, preview.Diagnostics)
require.Equal(t, "group", preview.Parameters[0].Name)
coderd/users_test.go
+2 -2
View File
@@ -117,8 +117,8 @@ func TestFirstUser(t *testing.T) {
_, err := client.CreateFirstUser(ctx, req)
require.NoError(t, err)
_ = testutil.RequireRecvCtx(ctx, t, trialGenerated)
_ = testutil.RequireRecvCtx(ctx, t, entitlementsRefreshed)
_ = testutil.TryReceive(ctx, t, trialGenerated)
_ = testutil.TryReceive(ctx, t, entitlementsRefreshed)
})
}
coderd/workspaceagents_test.go
+11 -11
View File
@@ -653,7 +653,7 @@ func TestWorkspaceAgentClientCoordinate_ResumeToken(t *testing.T) {
// random value.
originalResumeToken, err := connectToCoordinatorAndFetchResumeToken(ctx, logger, client, agentAndBuild.WorkspaceAgent.ID, "")
require.NoError(t, err)
originalPeerID := testutil.RequireRecvCtx(ctx, t, resumeTokenProvider.generateCalls)
originalPeerID := testutil.TryReceive(ctx, t, resumeTokenProvider.generateCalls)
require.NotEqual(t, originalPeerID, uuid.Nil)
// Connect with a valid resume token, and ensure that the peer ID is set to
@@ -661,9 +661,9 @@ func TestWorkspaceAgentClientCoordinate_ResumeToken(t *testing.T) {
clock.Advance(time.Second)
newResumeToken, err := connectToCoordinatorAndFetchResumeToken(ctx, logger, client, agentAndBuild.WorkspaceAgent.ID, originalResumeToken)
require.NoError(t, err)
verifiedToken := testutil.RequireRecvCtx(ctx, t, resumeTokenProvider.verifyCalls)
verifiedToken := testutil.TryReceive(ctx, t, resumeTokenProvider.verifyCalls)
require.Equal(t, originalResumeToken, verifiedToken)
newPeerID := testutil.RequireRecvCtx(ctx, t, resumeTokenProvider.generateCalls)
newPeerID := testutil.TryReceive(ctx, t, resumeTokenProvider.generateCalls)
require.Equal(t, originalPeerID, newPeerID)
require.NotEqual(t, originalResumeToken, newResumeToken)
@@ -677,7 +677,7 @@ func TestWorkspaceAgentClientCoordinate_ResumeToken(t *testing.T) {
require.Equal(t, http.StatusUnauthorized, sdkErr.StatusCode())
require.Len(t, sdkErr.Validations, 1)
require.Equal(t, "resume_token", sdkErr.Validations[0].Field)
verifiedToken = testutil.RequireRecvCtx(ctx, t, resumeTokenProvider.verifyCalls)
verifiedToken = testutil.TryReceive(ctx, t, resumeTokenProvider.verifyCalls)
require.Equal(t, "invalid", verifiedToken)
select {
@@ -725,7 +725,7 @@ func TestWorkspaceAgentClientCoordinate_ResumeToken(t *testing.T) {
// random value.
originalResumeToken, err := connectToCoordinatorAndFetchResumeToken(ctx, logger, client, agentAndBuild.WorkspaceAgent.ID, "")
require.NoError(t, err)
originalPeerID := testutil.RequireRecvCtx(ctx, t, resumeTokenProvider.generateCalls)
originalPeerID := testutil.TryReceive(ctx, t, resumeTokenProvider.generateCalls)
require.NotEqual(t, originalPeerID, uuid.Nil)
// Connect with an outdated token, and ensure that the peer ID is set to a
@@ -739,9 +739,9 @@ func TestWorkspaceAgentClientCoordinate_ResumeToken(t *testing.T) {
clock.Advance(time.Second)
newResumeToken, err := connectToCoordinatorAndFetchResumeToken(ctx, logger, client, agentAndBuild.WorkspaceAgent.ID, outdatedToken)
require.NoError(t, err)
verifiedToken := testutil.RequireRecvCtx(ctx, t, resumeTokenProvider.verifyCalls)
verifiedToken := testutil.TryReceive(ctx, t, resumeTokenProvider.verifyCalls)
require.Equal(t, outdatedToken, verifiedToken)
newPeerID := testutil.RequireRecvCtx(ctx, t, resumeTokenProvider.generateCalls)
newPeerID := testutil.TryReceive(ctx, t, resumeTokenProvider.generateCalls)
require.NotEqual(t, originalPeerID, newPeerID)
require.NotEqual(t, originalResumeToken, newResumeToken)
})
@@ -1912,8 +1912,8 @@ func TestWorkspaceAgent_Metadata_CatchMemoryLeak(t *testing.T) {
// testing it is not straightforward.
db.err.Store(&wantErr)
testutil.RequireRecvCtx(ctx, t, metadataDone)
testutil.RequireRecvCtx(ctx, t, postDone)
testutil.TryReceive(ctx, t, metadataDone)
testutil.TryReceive(ctx, t, postDone)
}
func TestWorkspaceAgent_Startup(t *testing.T) {
@@ -2358,7 +2358,7 @@ func TestUserTailnetTelemetry(t *testing.T) {
defer wsConn.Close(websocket.StatusNormalClosure, "done")
// Check telemetry
snapshot := testutil.RequireRecvCtx(ctx, t, fTelemetry.snapshots)
snapshot := testutil.TryReceive(ctx, t, fTelemetry.snapshots)
require.Len(t, snapshot.UserTailnetConnections, 1)
telemetryConnection := snapshot.UserTailnetConnections[0]
require.Equal(t, memberUser.ID.String(), telemetryConnection.UserID)
@@ -2373,7 +2373,7 @@ func TestUserTailnetTelemetry(t *testing.T) {
err = wsConn.Close(websocket.StatusNormalClosure, "done")
require.NoError(t, err)
snapshot = testutil.RequireRecvCtx(ctx, t, fTelemetry.snapshots)
snapshot = testutil.TryReceive(ctx, t, fTelemetry.snapshots)
require.Len(t, snapshot.UserTailnetConnections, 1)
telemetryDisconnection := snapshot.UserTailnetConnections[0]
require.Equal(t, memberUser.ID.String(), telemetryDisconnection.UserID)
coderd/workspaceagentsrpc_internal_test.go
+2 -2
View File
@@ -90,7 +90,7 @@ func TestAgentConnectionMonitor_ContextCancel(t *testing.T) {
fConn.requireEventuallyClosed(t, websocket.StatusGoingAway, "canceled")
// make sure we got at least one additional update on close
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
m := fUpdater.getUpdates()
require.Greater(t, m, n)
}
@@ -293,7 +293,7 @@ func TestAgentConnectionMonitor_StartClose(t *testing.T) {
uut.close()
close(closed)
}()
_ = testutil.RequireRecvCtx(ctx, t, closed)
_ = testutil.TryReceive(ctx, t, closed)
}
type fakePingerCloser struct {
coderd/workspaceupdates_test.go
+4 -4
View File
@@ -108,7 +108,7 @@ func TestWorkspaceUpdates(t *testing.T) {
_ = sub.Close()
})
update := testutil.RequireRecvCtx(ctx, t, sub.Updates())
update := testutil.TryReceive(ctx, t, sub.Updates())
slices.SortFunc(update.UpsertedWorkspaces, func(a, b *proto.Workspace) int {
return strings.Compare(a.Name, b.Name)
})
@@ -185,7 +185,7 @@ func TestWorkspaceUpdates(t *testing.T) {
WorkspaceID: ws1ID,
})
update = testutil.RequireRecvCtx(ctx, t, sub.Updates())
update = testutil.TryReceive(ctx, t, sub.Updates())
slices.SortFunc(update.UpsertedWorkspaces, func(a, b *proto.Workspace) int {
return strings.Compare(a.Name, b.Name)
})
@@ -284,7 +284,7 @@ func TestWorkspaceUpdates(t *testing.T) {
DeletedAgents: []*proto.Agent{},
}
update := testutil.RequireRecvCtx(ctx, t, sub.Updates())
update := testutil.TryReceive(ctx, t, sub.Updates())
slices.SortFunc(update.UpsertedWorkspaces, func(a, b *proto.Workspace) int {
return strings.Compare(a.Name, b.Name)
})
@@ -296,7 +296,7 @@ func TestWorkspaceUpdates(t *testing.T) {
_ = resub.Close()
})
update = testutil.RequireRecvCtx(ctx, t, resub.Updates())
update = testutil.TryReceive(ctx, t, resub.Updates())
slices.SortFunc(update.UpsertedWorkspaces, func(a, b *proto.Workspace) int {
return strings.Compare(a.Name, b.Name)
})
codersdk/agentsdk/logs_internal_test.go
+30 -30
View File
@@ -63,10 +63,10 @@ func TestLogSender_Mainline(t *testing.T) {
// since neither source has even been flushed, it should immediately Flush
// both, although the order is not controlled
var logReqs []*proto.BatchCreateLogsRequest
logReqs = append(logReqs, testutil.RequireRecvCtx(ctx, t, fDest.reqs))
testutil.RequireSendCtx(ctx, t, fDest.resps, &proto.BatchCreateLogsResponse{})
logReqs = append(logReqs, testutil.RequireRecvCtx(ctx, t, fDest.reqs))
testutil.RequireSendCtx(ctx, t, fDest.resps, &proto.BatchCreateLogsResponse{})
logReqs = append(logReqs, testutil.TryReceive(ctx, t, fDest.reqs))
testutil.RequireSend(ctx, t, fDest.resps, &proto.BatchCreateLogsResponse{})
logReqs = append(logReqs, testutil.TryReceive(ctx, t, fDest.reqs))
testutil.RequireSend(ctx, t, fDest.resps, &proto.BatchCreateLogsResponse{})
for _, req := range logReqs {
require.NotNil(t, req)
srcID, err := uuid.FromBytes(req.LogSourceId)
@@ -98,8 +98,8 @@ func TestLogSender_Mainline(t *testing.T) {
})
uut.Flush(ls1)
req := testutil.RequireRecvCtx(ctx, t, fDest.reqs)
testutil.RequireSendCtx(ctx, t, fDest.resps, &proto.BatchCreateLogsResponse{})
req := testutil.TryReceive(ctx, t, fDest.reqs)
testutil.RequireSend(ctx, t, fDest.resps, &proto.BatchCreateLogsResponse{})
// give ourselves a 25% buffer if we're right on the cusp of a tick
require.LessOrEqual(t, time.Since(t1), flushInterval*5/4)
require.NotNil(t, req)
@@ -108,11 +108,11 @@ func TestLogSender_Mainline(t *testing.T) {
require.Equal(t, proto.Log_DEBUG, req.Logs[0].GetLevel())
require.Equal(t, t1, req.Logs[0].GetCreatedAt().AsTime())
err := testutil.RequireRecvCtx(ctx, t, empty)
err := testutil.TryReceive(ctx, t, empty)
require.NoError(t, err)
cancel()
err = testutil.RequireRecvCtx(testCtx, t, loopErr)
err = testutil.TryReceive(testCtx, t, loopErr)
require.ErrorIs(t, err, context.Canceled)
// we can still enqueue more logs after SendLoop returns
@@ -151,16 +151,16 @@ func TestLogSender_LogLimitExceeded(t *testing.T) {
loopErr <- err
}()
req := testutil.RequireRecvCtx(ctx, t, fDest.reqs)
req := testutil.TryReceive(ctx, t, fDest.reqs)
require.NotNil(t, req)
testutil.RequireSendCtx(ctx, t, fDest.resps,
testutil.RequireSend(ctx, t, fDest.resps,
&proto.BatchCreateLogsResponse{LogLimitExceeded: true})
err := testutil.RequireRecvCtx(ctx, t, loopErr)
err := testutil.TryReceive(ctx, t, loopErr)
require.ErrorIs(t, err, ErrLogLimitExceeded)
// Should also unblock WaitUntilEmpty
err = testutil.RequireRecvCtx(ctx, t, empty)
err = testutil.TryReceive(ctx, t, empty)
require.NoError(t, err)
// we can still enqueue more logs after SendLoop returns, but they don't
@@ -179,7 +179,7 @@ func TestLogSender_LogLimitExceeded(t *testing.T) {
err := uut.SendLoop(ctx, fDest)
loopErr <- err
}()
err = testutil.RequireRecvCtx(ctx, t, loopErr)
err = testutil.TryReceive(ctx, t, loopErr)
require.ErrorIs(t, err, ErrLogLimitExceeded)
}
@@ -217,15 +217,15 @@ func TestLogSender_SkipHugeLog(t *testing.T) {
loopErr <- err
}()
req := testutil.RequireRecvCtx(ctx, t, fDest.reqs)
req := testutil.TryReceive(ctx, t, fDest.reqs)
require.NotNil(t, req)
require.Len(t, req.Logs, 1, "it should skip the huge log")
require.Equal(t, "test log 1, src 1", req.Logs[0].GetOutput())
require.Equal(t, proto.Log_INFO, req.Logs[0].GetLevel())
testutil.RequireSendCtx(ctx, t, fDest.resps, &proto.BatchCreateLogsResponse{})
testutil.RequireSend(ctx, t, fDest.resps, &proto.BatchCreateLogsResponse{})
cancel()
err := testutil.RequireRecvCtx(testCtx, t, loopErr)
err := testutil.TryReceive(testCtx, t, loopErr)
require.ErrorIs(t, err, context.Canceled)
}
@@ -258,7 +258,7 @@ func TestLogSender_InvalidUTF8(t *testing.T) {
loopErr <- err
}()
req := testutil.RequireRecvCtx(ctx, t, fDest.reqs)
req := testutil.TryReceive(ctx, t, fDest.reqs)
require.NotNil(t, req)
require.Len(t, req.Logs, 2, "it should sanitize invalid UTF-8, but still send")
// the 0xc3, 0x28 is an invalid 2-byte sequence in UTF-8. The sanitizer replaces 0xc3 with ❌, and then
@@ -267,10 +267,10 @@ func TestLogSender_InvalidUTF8(t *testing.T) {
require.Equal(t, proto.Log_INFO, req.Logs[0].GetLevel())
require.Equal(t, "test log 1, src 1", req.Logs[1].GetOutput())
require.Equal(t, proto.Log_INFO, req.Logs[1].GetLevel())
testutil.RequireSendCtx(ctx, t, fDest.resps, &proto.BatchCreateLogsResponse{})
testutil.RequireSend(ctx, t, fDest.resps, &proto.BatchCreateLogsResponse{})
cancel()
err := testutil.RequireRecvCtx(testCtx, t, loopErr)
err := testutil.TryReceive(testCtx, t, loopErr)
require.ErrorIs(t, err, context.Canceled)
}
@@ -303,24 +303,24 @@ func TestLogSender_Batch(t *testing.T) {
// with 60k logs, we should split into two updates to avoid going over 1MiB, since each log
// is about 21 bytes.
gotLogs := 0
req := testutil.RequireRecvCtx(ctx, t, fDest.reqs)
req := testutil.TryReceive(ctx, t, fDest.reqs)
require.NotNil(t, req)
gotLogs += len(req.Logs)
wire, err := protobuf.Marshal(req)
require.NoError(t, err)
require.Less(t, len(wire), maxBytesPerBatch, "wire should not exceed 1MiB")
testutil.RequireSendCtx(ctx, t, fDest.resps, &proto.BatchCreateLogsResponse{})
req = testutil.RequireRecvCtx(ctx, t, fDest.reqs)
testutil.RequireSend(ctx, t, fDest.resps, &proto.BatchCreateLogsResponse{})
req = testutil.TryReceive(ctx, t, fDest.reqs)
require.NotNil(t, req)
gotLogs += len(req.Logs)
wire, err = protobuf.Marshal(req)
require.NoError(t, err)
require.Less(t, len(wire), maxBytesPerBatch, "wire should not exceed 1MiB")
require.Equal(t, 60000, gotLogs)
testutil.RequireSendCtx(ctx, t, fDest.resps, &proto.BatchCreateLogsResponse{})
testutil.RequireSend(ctx, t, fDest.resps, &proto.BatchCreateLogsResponse{})
cancel()
err = testutil.RequireRecvCtx(testCtx, t, loopErr)
err = testutil.TryReceive(testCtx, t, loopErr)
require.ErrorIs(t, err, context.Canceled)
}
@@ -367,12 +367,12 @@ func TestLogSender_MaxQueuedLogs(t *testing.T) {
// #1 come in 2 updates, plus 1 update for source #2.
logsBySource := make(map[uuid.UUID]int)
for i := 0; i < 3; i++ {
req := testutil.RequireRecvCtx(ctx, t, fDest.reqs)
req := testutil.TryReceive(ctx, t, fDest.reqs)
require.NotNil(t, req)
srcID, err := uuid.FromBytes(req.LogSourceId)
require.NoError(t, err)
logsBySource[srcID] += len(req.Logs)
testutil.RequireSendCtx(ctx, t, fDest.resps, &proto.BatchCreateLogsResponse{})
testutil.RequireSend(ctx, t, fDest.resps, &proto.BatchCreateLogsResponse{})
}
require.Equal(t, map[uuid.UUID]int{
ls1: n,
@@ -380,7 +380,7 @@ func TestLogSender_MaxQueuedLogs(t *testing.T) {
}, logsBySource)
cancel()
err := testutil.RequireRecvCtx(testCtx, t, loopErr)
err := testutil.TryReceive(testCtx, t, loopErr)
require.ErrorIs(t, err, context.Canceled)
}
@@ -408,10 +408,10 @@ func TestLogSender_SendError(t *testing.T) {
loopErr <- err
}()
req := testutil.RequireRecvCtx(ctx, t, fDest.reqs)
req := testutil.TryReceive(ctx, t, fDest.reqs)
require.NotNil(t, req)
err := testutil.RequireRecvCtx(ctx, t, loopErr)
err := testutil.TryReceive(ctx, t, loopErr)
require.ErrorIs(t, err, expectedErr)
// we can still enqueue more logs after SendLoop returns
@@ -448,7 +448,7 @@ func TestLogSender_WaitUntilEmpty_ContextExpired(t *testing.T) {
}()
cancel()
err := testutil.RequireRecvCtx(testCtx, t, empty)
err := testutil.TryReceive(testCtx, t, empty)
require.ErrorIs(t, err, context.Canceled)
}
codersdk/workspacesdk/dialer_test.go
+16 -16
View File
@@ -80,15 +80,15 @@ func TestWebsocketDialer_TokenController(t *testing.T) {
clientCh <- clients
}()
call := testutil.RequireRecvCtx(ctx, t, fTokenProv.tokenCalls)
call := testutil.TryReceive(ctx, t, fTokenProv.tokenCalls)
call <- tokenResponse{"test token", true}
gotToken := <-dialTokens
require.Equal(t, "test token", gotToken)
clients := testutil.RequireRecvCtx(ctx, t, clientCh)
clients := testutil.TryReceive(ctx, t, clientCh)
clients.Closer.Close()
err = testutil.RequireRecvCtx(ctx, t, wsErr)
err = testutil.TryReceive(ctx, t, wsErr)
require.NoError(t, err)
clientCh = make(chan tailnet.ControlProtocolClients, 1)
@@ -98,16 +98,16 @@ func TestWebsocketDialer_TokenController(t *testing.T) {
clientCh <- clients
}()
call = testutil.RequireRecvCtx(ctx, t, fTokenProv.tokenCalls)
call = testutil.TryReceive(ctx, t, fTokenProv.tokenCalls)
call <- tokenResponse{"test token", false}
gotToken = <-dialTokens
require.Equal(t, "", gotToken)
clients = testutil.RequireRecvCtx(ctx, t, clientCh)
clients = testutil.TryReceive(ctx, t, clientCh)
require.Nil(t, clients.WorkspaceUpdates)
clients.Closer.Close()
err = testutil.RequireRecvCtx(ctx, t, wsErr)
err = testutil.TryReceive(ctx, t, wsErr)
require.NoError(t, err)
}
@@ -165,10 +165,10 @@ func TestWebsocketDialer_NoTokenController(t *testing.T) {
gotToken := <-dialTokens
require.Equal(t, "", gotToken)
clients := testutil.RequireRecvCtx(ctx, t, clientCh)
clients := testutil.TryReceive(ctx, t, clientCh)
clients.Closer.Close()
err = testutil.RequireRecvCtx(ctx, t, wsErr)
err = testutil.TryReceive(ctx, t, wsErr)
require.NoError(t, err)
}
@@ -233,12 +233,12 @@ func TestWebsocketDialer_ResumeTokenFailure(t *testing.T) {
errCh <- err
}()
call := testutil.RequireRecvCtx(ctx, t, fTokenProv.tokenCalls)
call := testutil.TryReceive(ctx, t, fTokenProv.tokenCalls)
call <- tokenResponse{"test token", true}
gotToken := <-dialTokens
require.Equal(t, "test token", gotToken)
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
require.Error(t, err)
// redial should not use the token
@@ -251,10 +251,10 @@ func TestWebsocketDialer_ResumeTokenFailure(t *testing.T) {
gotToken = <-dialTokens
require.Equal(t, "", gotToken)
clients := testutil.RequireRecvCtx(ctx, t, clientCh)
clients := testutil.TryReceive(ctx, t, clientCh)
require.Error(t, err)
clients.Closer.Close()
err = testutil.RequireRecvCtx(ctx, t, wsErr)
err = testutil.TryReceive(ctx, t, wsErr)
require.NoError(t, err)
// Successful dial should reset to using token again
@@ -262,11 +262,11 @@ func TestWebsocketDialer_ResumeTokenFailure(t *testing.T) {
_, err := uut.Dial(ctx, fTokenProv)
errCh <- err
}()
call = testutil.RequireRecvCtx(ctx, t, fTokenProv.tokenCalls)
call = testutil.TryReceive(ctx, t, fTokenProv.tokenCalls)
call <- tokenResponse{"test token", true}
gotToken = <-dialTokens
require.Equal(t, "test token", gotToken)
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
require.Error(t, err)
}
@@ -305,7 +305,7 @@ func TestWebsocketDialer_UplevelVersion(t *testing.T) {
errCh <- err
}()
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
var sdkErr *codersdk.Error
require.ErrorAs(t, err, &sdkErr)
require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode())
@@ -387,7 +387,7 @@ func TestWebsocketDialer_WorkspaceUpdates(t *testing.T) {
clients.Closer.Close()
err = testutil.RequireRecvCtx(ctx, t, wsErr)
err = testutil.TryReceive(ctx, t, wsErr)
require.NoError(t, err)
}
enterprise/tailnet/pgcoord_internal_test.go
+1 -1
View File
@@ -427,7 +427,7 @@ func TestPGCoordinatorUnhealthy(t *testing.T) {
pID := uuid.UUID{5}
_, resps := coordinator.Coordinate(ctx, pID, "test", agpl.AgentCoordinateeAuth{ID: pID})
resp := testutil.RequireRecvCtx(ctx, t, resps)
resp := testutil.TryReceive(ctx, t, resps)
require.Nil(t, resp, "channel should be closed")
// give the coordinator some time to process any pending work. We are
enterprise/tailnet/pgcoord_test.go
+2 -2
View File
@@ -943,9 +943,9 @@ func TestPGCoordinatorPropogatedPeerContext(t *testing.T) {
reqs, _ := c1.Coordinate(peerCtx, peerID, "peer1", auth)
testutil.RequireSendCtx(ctx, t, reqs, &proto.CoordinateRequest{AddTunnel: &proto.CoordinateRequest_Tunnel{Id: agpl.UUIDToByteSlice(agentID)}})
testutil.RequireSend(ctx, t, reqs, &proto.CoordinateRequest{AddTunnel: &proto.CoordinateRequest_Tunnel{Id: agpl.UUIDToByteSlice(agentID)}})
_ = testutil.RequireRecvCtx(ctx, t, ch)
_ = testutil.TryReceive(ctx, t, ch)
}
func assertEventuallyStatus(ctx context.Context, t *testing.T, store database.Store, agentID uuid.UUID, status database.TailnetStatus) {
enterprise/wsproxy/wsproxy_test.go
+3 -3
View File
@@ -780,7 +780,7 @@ func TestWorkspaceProxyDERPMeshProbe(t *testing.T) {
require.NoError(t, err, "failed to force proxy to re-register")
// Wait for the ping to fail.
replicaErr := testutil.RequireRecvCtx(ctx, t, replicaPingErr)
replicaErr := testutil.TryReceive(ctx, t, replicaPingErr)
require.NotEmpty(t, replicaErr, "replica ping error")
// GET /healthz-report
@@ -858,7 +858,7 @@ func TestWorkspaceProxyDERPMeshProbe(t *testing.T) {
// Wait for the ping to fail.
for {
replicaErr := testutil.RequireRecvCtx(ctx, t, replicaPingErr)
replicaErr := testutil.TryReceive(ctx, t, replicaPingErr)
t.Log("replica ping error:", replicaErr)
if replicaErr != "" {
break
@@ -892,7 +892,7 @@ func TestWorkspaceProxyDERPMeshProbe(t *testing.T) {
// Wait for the ping to be skipped.
for {
replicaErr := testutil.RequireRecvCtx(ctx, t, replicaPingErr)
replicaErr := testutil.TryReceive(ctx, t, replicaPingErr)
t.Log("replica ping error:", replicaErr)
// Should be empty because there are no more peers. This was where
// the regression was.
scaletest/createworkspaces/run_test.go
+1 -1
View File
@@ -293,7 +293,7 @@ func Test_Runner(t *testing.T) {
<-done
t.Log("canceled scaletest workspace creation")
// Ensure we have a job to interrogate
runningJob := testutil.RequireRecvCtx(testutil.Context(t, testutil.WaitShort), t, jobCh)
runningJob := testutil.TryReceive(testutil.Context(t, testutil.WaitShort), t, jobCh)
require.NotZero(t, runningJob.ID)
// When we run the cleanup, it should be canceled
tailnet/configmaps_internal_test.go
+68 -68
View File
@@ -40,7 +40,7 @@ func TestConfigMaps_setAddresses_different(t *testing.T) {
addrs := []netip.Prefix{netip.MustParsePrefix("192.168.0.200/32")}
uut.setAddresses(addrs)
nm := testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
nm := testutil.TryReceive(ctx, t, fEng.setNetworkMap)
require.Equal(t, addrs, nm.Addresses)
// here were in the middle of a reconfig, blocked on a channel write to fEng.reconfig
@@ -55,22 +55,22 @@ func TestConfigMaps_setAddresses_different(t *testing.T) {
}
uut.setAddresses(addrs2)
r := testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
r := testutil.TryReceive(ctx, t, fEng.reconfig)
require.Equal(t, addrs, r.wg.Addresses)
require.Equal(t, addrs, r.router.LocalAddrs)
f := testutil.RequireRecvCtx(ctx, t, fEng.filter)
f := testutil.TryReceive(ctx, t, fEng.filter)
fr := f.CheckTCP(netip.MustParseAddr("33.44.55.66"), netip.MustParseAddr("192.168.0.200"), 5555)
require.Equal(t, filter.Accept, fr)
fr = f.CheckTCP(netip.MustParseAddr("33.44.55.66"), netip.MustParseAddr("10.20.30.40"), 5555)
require.Equal(t, filter.Drop, fr, "first addr config should not include 10.20.30.40")
// we should get another round of configurations from the second set of addrs
nm = testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
nm = testutil.TryReceive(ctx, t, fEng.setNetworkMap)
require.Equal(t, addrs2, nm.Addresses)
r = testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
r = testutil.TryReceive(ctx, t, fEng.reconfig)
require.Equal(t, addrs2, r.wg.Addresses)
require.Equal(t, addrs2, r.router.LocalAddrs)
f = testutil.RequireRecvCtx(ctx, t, fEng.filter)
f = testutil.TryReceive(ctx, t, fEng.filter)
fr = f.CheckTCP(netip.MustParseAddr("33.44.55.66"), netip.MustParseAddr("192.168.0.200"), 5555)
require.Equal(t, filter.Accept, fr)
fr = f.CheckTCP(netip.MustParseAddr("33.44.55.66"), netip.MustParseAddr("10.20.30.40"), 5555)
@@ -81,7 +81,7 @@ func TestConfigMaps_setAddresses_different(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestConfigMaps_setAddresses_same(t *testing.T) {
@@ -112,7 +112,7 @@ func TestConfigMaps_setAddresses_same(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestConfigMaps_updatePeers_new(t *testing.T) {
@@ -160,8 +160,8 @@ func TestConfigMaps_updatePeers_new(t *testing.T) {
}
uut.updatePeers(updates)
nm := testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
r := testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
nm := testutil.TryReceive(ctx, t, fEng.setNetworkMap)
r := testutil.TryReceive(ctx, t, fEng.reconfig)
require.Len(t, nm.Peers, 2)
n1 := getNodeWithID(t, nm.Peers, 1)
@@ -182,7 +182,7 @@ func TestConfigMaps_updatePeers_new(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestConfigMaps_updatePeers_new_waitForHandshake_neverConfigures(t *testing.T) {
@@ -226,7 +226,7 @@ func TestConfigMaps_updatePeers_new_waitForHandshake_neverConfigures(t *testing.
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestConfigMaps_updatePeers_new_waitForHandshake_outOfOrder(t *testing.T) {
@@ -279,8 +279,8 @@ func TestConfigMaps_updatePeers_new_waitForHandshake_outOfOrder(t *testing.T) {
// it should now send the peer to the netmap
nm := testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
r := testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
nm := testutil.TryReceive(ctx, t, fEng.setNetworkMap)
r := testutil.TryReceive(ctx, t, fEng.reconfig)
require.Len(t, nm.Peers, 1)
n1 := getNodeWithID(t, nm.Peers, 1)
@@ -297,7 +297,7 @@ func TestConfigMaps_updatePeers_new_waitForHandshake_outOfOrder(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestConfigMaps_updatePeers_new_waitForHandshake(t *testing.T) {
@@ -350,8 +350,8 @@ func TestConfigMaps_updatePeers_new_waitForHandshake(t *testing.T) {
// it should now send the peer to the netmap
nm := testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
r := testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
nm := testutil.TryReceive(ctx, t, fEng.setNetworkMap)
r := testutil.TryReceive(ctx, t, fEng.reconfig)
require.Len(t, nm.Peers, 1)
n1 := getNodeWithID(t, nm.Peers, 1)
@@ -368,7 +368,7 @@ func TestConfigMaps_updatePeers_new_waitForHandshake(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestConfigMaps_updatePeers_new_waitForHandshake_timeout(t *testing.T) {
@@ -408,8 +408,8 @@ func TestConfigMaps_updatePeers_new_waitForHandshake_timeout(t *testing.T) {
// it should now send the peer to the netmap
nm := testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
r := testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
nm := testutil.TryReceive(ctx, t, fEng.setNetworkMap)
r := testutil.TryReceive(ctx, t, fEng.reconfig)
require.Len(t, nm.Peers, 1)
n1 := getNodeWithID(t, nm.Peers, 1)
@@ -426,7 +426,7 @@ func TestConfigMaps_updatePeers_new_waitForHandshake_timeout(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestConfigMaps_updatePeers_same(t *testing.T) {
@@ -485,7 +485,7 @@ func TestConfigMaps_updatePeers_same(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestConfigMaps_updatePeers_disconnect(t *testing.T) {
@@ -543,8 +543,8 @@ func TestConfigMaps_updatePeers_disconnect(t *testing.T) {
assert.False(t, timer.Stop(), "timer was not stopped")
// Then, configure engine without the peer.
nm := testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
r := testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
nm := testutil.TryReceive(ctx, t, fEng.setNetworkMap)
r := testutil.TryReceive(ctx, t, fEng.reconfig)
require.Len(t, nm.Peers, 0)
require.Len(t, r.wg.Peers, 0)
@@ -553,7 +553,7 @@ func TestConfigMaps_updatePeers_disconnect(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestConfigMaps_updatePeers_lost(t *testing.T) {
@@ -585,11 +585,11 @@ func TestConfigMaps_updatePeers_lost(t *testing.T) {
},
}
uut.updatePeers(updates)
nm := testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
r := testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
nm := testutil.TryReceive(ctx, t, fEng.setNetworkMap)
r := testutil.TryReceive(ctx, t, fEng.reconfig)
require.Len(t, nm.Peers, 1)
require.Len(t, r.wg.Peers, 1)
_ = testutil.RequireRecvCtx(ctx, t, s1)
_ = testutil.TryReceive(ctx, t, s1)
mClock.Advance(5 * time.Second).MustWait(ctx)
@@ -598,7 +598,7 @@ func TestConfigMaps_updatePeers_lost(t *testing.T) {
updates[0].Kind = proto.CoordinateResponse_PeerUpdate_LOST
updates[0].Node = nil
uut.updatePeers(updates)
_ = testutil.RequireRecvCtx(ctx, t, s2)
_ = testutil.TryReceive(ctx, t, s2)
// No reprogramming yet, since we keep the peer around.
select {
@@ -614,7 +614,7 @@ func TestConfigMaps_updatePeers_lost(t *testing.T) {
s3 := expectStatusWithHandshake(ctx, t, fEng, p1Node.Key, lh)
// 5 seconds have already elapsed from above
mClock.Advance(lostTimeout - 5*time.Second).MustWait(ctx)
_ = testutil.RequireRecvCtx(ctx, t, s3)
_ = testutil.TryReceive(ctx, t, s3)
select {
case <-fEng.setNetworkMap:
t.Fatal("should not reprogram")
@@ -627,18 +627,18 @@ func TestConfigMaps_updatePeers_lost(t *testing.T) {
s4 := expectStatusWithHandshake(ctx, t, fEng, p1Node.Key, lh)
mClock.Advance(time.Minute).MustWait(ctx)
nm = testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
r = testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
nm = testutil.TryReceive(ctx, t, fEng.setNetworkMap)
r = testutil.TryReceive(ctx, t, fEng.reconfig)
require.Len(t, nm.Peers, 0)
require.Len(t, r.wg.Peers, 0)
_ = testutil.RequireRecvCtx(ctx, t, s4)
_ = testutil.TryReceive(ctx, t, s4)
done := make(chan struct{})
go func() {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestConfigMaps_updatePeers_lost_and_found(t *testing.T) {
@@ -670,11 +670,11 @@ func TestConfigMaps_updatePeers_lost_and_found(t *testing.T) {
},
}
uut.updatePeers(updates)
nm := testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
r := testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
nm := testutil.TryReceive(ctx, t, fEng.setNetworkMap)
r := testutil.TryReceive(ctx, t, fEng.reconfig)
require.Len(t, nm.Peers, 1)
require.Len(t, r.wg.Peers, 1)
_ = testutil.RequireRecvCtx(ctx, t, s1)
_ = testutil.TryReceive(ctx, t, s1)
mClock.Advance(5 * time.Second).MustWait(ctx)
@@ -683,7 +683,7 @@ func TestConfigMaps_updatePeers_lost_and_found(t *testing.T) {
updates[0].Kind = proto.CoordinateResponse_PeerUpdate_LOST
updates[0].Node = nil
uut.updatePeers(updates)
_ = testutil.RequireRecvCtx(ctx, t, s2)
_ = testutil.TryReceive(ctx, t, s2)
// No reprogramming yet, since we keep the peer around.
select {
@@ -699,7 +699,7 @@ func TestConfigMaps_updatePeers_lost_and_found(t *testing.T) {
updates[0].Kind = proto.CoordinateResponse_PeerUpdate_NODE
updates[0].Node = p1n
uut.updatePeers(updates)
_ = testutil.RequireRecvCtx(ctx, t, s3)
_ = testutil.TryReceive(ctx, t, s3)
// This does not trigger reprogramming, because we never removed the node
select {
case <-fEng.setNetworkMap:
@@ -723,7 +723,7 @@ func TestConfigMaps_updatePeers_lost_and_found(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestConfigMaps_setAllPeersLost(t *testing.T) {
@@ -764,11 +764,11 @@ func TestConfigMaps_setAllPeersLost(t *testing.T) {
},
}
uut.updatePeers(updates)
nm := testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
r := testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
nm := testutil.TryReceive(ctx, t, fEng.setNetworkMap)
r := testutil.TryReceive(ctx, t, fEng.reconfig)
require.Len(t, nm.Peers, 2)
require.Len(t, r.wg.Peers, 2)
_ = testutil.RequireRecvCtx(ctx, t, s1)
_ = testutil.TryReceive(ctx, t, s1)
mClock.Advance(5 * time.Second).MustWait(ctx)
uut.setAllPeersLost()
@@ -787,20 +787,20 @@ func TestConfigMaps_setAllPeersLost(t *testing.T) {
d, w := mClock.AdvanceNext()
w.MustWait(ctx)
require.LessOrEqual(t, d, time.Millisecond)
_ = testutil.RequireRecvCtx(ctx, t, s2)
_ = testutil.TryReceive(ctx, t, s2)
nm = testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
r = testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
nm = testutil.TryReceive(ctx, t, fEng.setNetworkMap)
r = testutil.TryReceive(ctx, t, fEng.reconfig)
require.Len(t, nm.Peers, 1)
require.Len(t, r.wg.Peers, 1)
// Finally, advance the clock until after the timeout
s3 := expectStatusWithHandshake(ctx, t, fEng, p1Node.Key, start)
mClock.Advance(lostTimeout - d - 5*time.Second).MustWait(ctx)
_ = testutil.RequireRecvCtx(ctx, t, s3)
_ = testutil.TryReceive(ctx, t, s3)
nm = testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
r = testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
nm = testutil.TryReceive(ctx, t, fEng.setNetworkMap)
r = testutil.TryReceive(ctx, t, fEng.reconfig)
require.Len(t, nm.Peers, 0)
require.Len(t, r.wg.Peers, 0)
@@ -809,7 +809,7 @@ func TestConfigMaps_setAllPeersLost(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestConfigMaps_setBlockEndpoints_different(t *testing.T) {
@@ -842,8 +842,8 @@ func TestConfigMaps_setBlockEndpoints_different(t *testing.T) {
uut.setBlockEndpoints(true)
nm := testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
r := testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
nm := testutil.TryReceive(ctx, t, fEng.setNetworkMap)
r := testutil.TryReceive(ctx, t, fEng.reconfig)
require.Len(t, nm.Peers, 1)
require.Len(t, nm.Peers[0].Endpoints, 0)
require.Len(t, r.wg.Peers, 1)
@@ -853,7 +853,7 @@ func TestConfigMaps_setBlockEndpoints_different(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestConfigMaps_setBlockEndpoints_same(t *testing.T) {
@@ -896,7 +896,7 @@ func TestConfigMaps_setBlockEndpoints_same(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestConfigMaps_setDERPMap_different(t *testing.T) {
@@ -923,7 +923,7 @@ func TestConfigMaps_setDERPMap_different(t *testing.T) {
}
uut.setDERPMap(derpMap)
dm := testutil.RequireRecvCtx(ctx, t, fEng.setDERPMap)
dm := testutil.TryReceive(ctx, t, fEng.setDERPMap)
require.Len(t, dm.HomeParams.RegionScore, 1)
require.Equal(t, dm.HomeParams.RegionScore[1], 0.025)
require.Len(t, dm.Regions, 1)
@@ -937,7 +937,7 @@ func TestConfigMaps_setDERPMap_different(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestConfigMaps_setDERPMap_same(t *testing.T) {
@@ -1006,7 +1006,7 @@ func TestConfigMaps_setDERPMap_same(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestConfigMaps_fillPeerDiagnostics(t *testing.T) {
@@ -1066,7 +1066,7 @@ func TestConfigMaps_fillPeerDiagnostics(t *testing.T) {
// When: call fillPeerDiagnostics
d := PeerDiagnostics{DERPRegionNames: make(map[int]string)}
uut.fillPeerDiagnostics(&d, p1ID)
testutil.RequireRecvCtx(ctx, t, s0)
testutil.TryReceive(ctx, t, s0)
// Then:
require.Equal(t, map[int]string{1: "AUH", 1001: "DXB"}, d.DERPRegionNames)
@@ -1078,7 +1078,7 @@ func TestConfigMaps_fillPeerDiagnostics(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func expectStatusWithHandshake(
@@ -1152,7 +1152,7 @@ func TestConfigMaps_updatePeers_nonexist(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
})
}
}
@@ -1187,8 +1187,8 @@ func TestConfigMaps_addRemoveHosts(t *testing.T) {
})
// THEN: the engine is reconfigured with those same hosts
_ = testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
req := testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
_ = testutil.TryReceive(ctx, t, fEng.setNetworkMap)
req := testutil.TryReceive(ctx, t, fEng.reconfig)
require.Equal(t, req.dnsCfg, &dns.Config{
Routes: map[dnsname.FQDN][]*dnstype.Resolver{
suffix: nil,
@@ -1218,8 +1218,8 @@ func TestConfigMaps_addRemoveHosts(t *testing.T) {
})
// THEN: The engine is reconfigured with only the new hosts
_ = testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
req = testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
_ = testutil.TryReceive(ctx, t, fEng.setNetworkMap)
req = testutil.TryReceive(ctx, t, fEng.reconfig)
require.Equal(t, req.dnsCfg, &dns.Config{
Routes: map[dnsname.FQDN][]*dnstype.Resolver{
suffix: nil,
@@ -1237,8 +1237,8 @@ func TestConfigMaps_addRemoveHosts(t *testing.T) {
// WHEN: we remove all the hosts
uut.setHosts(map[dnsname.FQDN][]netip.Addr{})
_ = testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
req = testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
_ = testutil.TryReceive(ctx, t, fEng.setNetworkMap)
req = testutil.TryReceive(ctx, t, fEng.reconfig)
// THEN: the engine is reconfigured with an empty config
require.Equal(t, req.dnsCfg, &dns.Config{})
@@ -1248,7 +1248,7 @@ func TestConfigMaps_addRemoveHosts(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func newTestNode(id int) *Node {
@@ -1287,7 +1287,7 @@ func requireNeverConfigures(ctx context.Context, t *testing.T, uut *phased) {
}
assert.Equal(t, closed, uut.phase)
}()
_ = testutil.RequireRecvCtx(ctx, t, waiting)
_ = testutil.TryReceive(ctx, t, waiting)
}
type reconfigCall struct {
tailnet/conn_test.go
+4 -4
View File
@@ -79,7 +79,7 @@ func TestTailnet(t *testing.T) {
conn <- struct{}{}
}()
_ = testutil.RequireRecvCtx(ctx, t, listenDone)
_ = testutil.TryReceive(ctx, t, listenDone)
nc, err := w2.DialContextTCP(context.Background(), netip.AddrPortFrom(w1IP, 35565))
require.NoError(t, err)
_ = nc.Close()
@@ -92,7 +92,7 @@ func TestTailnet(t *testing.T) {
default:
}
})
node := testutil.RequireRecvCtx(ctx, t, nodes)
node := testutil.TryReceive(ctx, t, nodes)
// Ensure this connected over raw (not websocket) DERP!
require.Len(t, node.DERPForcedWebsocket, 0)
@@ -146,11 +146,11 @@ func TestTailnet(t *testing.T) {
_ = nc.Close()
}()
testutil.RequireRecvCtx(ctx, t, listening)
testutil.TryReceive(ctx, t, listening)
nc, err := w2.DialContextTCP(ctx, netip.AddrPortFrom(w1IP, 35565))
require.NoError(t, err)
_ = nc.Close()
testutil.RequireRecvCtx(ctx, t, done)
testutil.TryReceive(ctx, t, done)
nodes := make(chan *tailnet.Node, 1)
w2.SetNodeCallback(func(node *tailnet.Node) {
tailnet/controllers_test.go
+187 -187
View File
@@ -61,7 +61,7 @@ func TestInMemoryCoordination(t *testing.T) {
coordinationTest(ctx, t, uut, fConn, reqs, resps, agentID)
// Recv loop should be terminated by the server hanging up after Disconnect
err := testutil.RequireRecvCtx(ctx, t, uut.Wait())
err := testutil.TryReceive(ctx, t, uut.Wait())
require.ErrorIs(t, err, io.EOF)
}
@@ -118,7 +118,7 @@ func TestTunnelSrcCoordController_Mainline(t *testing.T) {
coordinationTest(ctx, t, uut, fConn, reqs, resps, agentID)
// Recv loop should be terminated by the server hanging up after Disconnect
err = testutil.RequireRecvCtx(ctx, t, uut.Wait())
err = testutil.TryReceive(ctx, t, uut.Wait())
require.ErrorIs(t, err, io.EOF)
}
@@ -147,22 +147,22 @@ func TestTunnelSrcCoordController_AddDestination(t *testing.T) {
// THEN: Controller sends AddTunnel for the destinations
for i := range 2 {
b0 := byte(i + 1)
call := testutil.RequireRecvCtx(ctx, t, client1.reqs)
call := testutil.TryReceive(ctx, t, client1.reqs)
require.Equal(t, b0, call.req.GetAddTunnel().GetId()[0])
testutil.RequireSendCtx(ctx, t, call.err, nil)
testutil.RequireSend(ctx, t, call.err, nil)
}
_ = testutil.RequireRecvCtx(ctx, t, addDone)
_ = testutil.TryReceive(ctx, t, addDone)
// THEN: Controller sets destinations on Coordinatee
require.Contains(t, fConn.tunnelDestinations, dest1)
require.Contains(t, fConn.tunnelDestinations, dest2)
// WHEN: Closed from server side and reconnects
respCall := testutil.RequireRecvCtx(ctx, t, client1.resps)
testutil.RequireSendCtx(ctx, t, respCall.err, io.EOF)
closeCall := testutil.RequireRecvCtx(ctx, t, client1.close)
testutil.RequireSendCtx(ctx, t, closeCall, nil)
err := testutil.RequireRecvCtx(ctx, t, cw1.Wait())
respCall := testutil.TryReceive(ctx, t, client1.resps)
testutil.RequireSend(ctx, t, respCall.err, io.EOF)
closeCall := testutil.TryReceive(ctx, t, client1.close)
testutil.RequireSend(ctx, t, closeCall, nil)
err := testutil.TryReceive(ctx, t, cw1.Wait())
require.ErrorIs(t, err, io.EOF)
client2 := newFakeCoordinatorClient(ctx, t)
cws := make(chan tailnet.CloserWaiter)
@@ -173,21 +173,21 @@ func TestTunnelSrcCoordController_AddDestination(t *testing.T) {
// THEN: should immediately send both destinations
var dests []byte
for range 2 {
call := testutil.RequireRecvCtx(ctx, t, client2.reqs)
call := testutil.TryReceive(ctx, t, client2.reqs)
dests = append(dests, call.req.GetAddTunnel().GetId()[0])
testutil.RequireSendCtx(ctx, t, call.err, nil)
testutil.RequireSend(ctx, t, call.err, nil)
}
slices.Sort(dests)
require.Equal(t, dests, []byte{1, 2})
cw2 := testutil.RequireRecvCtx(ctx, t, cws)
cw2 := testutil.TryReceive(ctx, t, cws)
// close client2
respCall = testutil.RequireRecvCtx(ctx, t, client2.resps)
testutil.RequireSendCtx(ctx, t, respCall.err, io.EOF)
closeCall = testutil.RequireRecvCtx(ctx, t, client2.close)
testutil.RequireSendCtx(ctx, t, closeCall, nil)
err = testutil.RequireRecvCtx(ctx, t, cw2.Wait())
respCall = testutil.TryReceive(ctx, t, client2.resps)
testutil.RequireSend(ctx, t, respCall.err, io.EOF)
closeCall = testutil.TryReceive(ctx, t, client2.close)
testutil.RequireSend(ctx, t, closeCall, nil)
err = testutil.TryReceive(ctx, t, cw2.Wait())
require.ErrorIs(t, err, io.EOF)
}
@@ -209,9 +209,9 @@ func TestTunnelSrcCoordController_RemoveDestination(t *testing.T) {
go func() {
cws <- uut.New(client1)
}()
call := testutil.RequireRecvCtx(ctx, t, client1.reqs)
testutil.RequireSendCtx(ctx, t, call.err, nil)
cw1 := testutil.RequireRecvCtx(ctx, t, cws)
call := testutil.TryReceive(ctx, t, client1.reqs)
testutil.RequireSend(ctx, t, call.err, nil)
cw1 := testutil.TryReceive(ctx, t, cws)
// WHEN: we remove one destination
removeDone := make(chan struct{})
@@ -221,17 +221,17 @@ func TestTunnelSrcCoordController_RemoveDestination(t *testing.T) {
}()
// THEN: Controller sends RemoveTunnel for the destination
call = testutil.RequireRecvCtx(ctx, t, client1.reqs)
call = testutil.TryReceive(ctx, t, client1.reqs)
require.Equal(t, dest1[:], call.req.GetRemoveTunnel().GetId())
testutil.RequireSendCtx(ctx, t, call.err, nil)
_ = testutil.RequireRecvCtx(ctx, t, removeDone)
testutil.RequireSend(ctx, t, call.err, nil)
_ = testutil.TryReceive(ctx, t, removeDone)
// WHEN: Closed from server side and reconnect
respCall := testutil.RequireRecvCtx(ctx, t, client1.resps)
testutil.RequireSendCtx(ctx, t, respCall.err, io.EOF)
closeCall := testutil.RequireRecvCtx(ctx, t, client1.close)
testutil.RequireSendCtx(ctx, t, closeCall, nil)
err := testutil.RequireRecvCtx(ctx, t, cw1.Wait())
respCall := testutil.TryReceive(ctx, t, client1.resps)
testutil.RequireSend(ctx, t, respCall.err, io.EOF)
closeCall := testutil.TryReceive(ctx, t, client1.close)
testutil.RequireSend(ctx, t, closeCall, nil)
err := testutil.TryReceive(ctx, t, cw1.Wait())
require.ErrorIs(t, err, io.EOF)
client2 := newFakeCoordinatorClient(ctx, t)
@@ -240,14 +240,14 @@ func TestTunnelSrcCoordController_RemoveDestination(t *testing.T) {
}()
// THEN: should immediately resolve without sending anything
cw2 := testutil.RequireRecvCtx(ctx, t, cws)
cw2 := testutil.TryReceive(ctx, t, cws)
// close client2
respCall = testutil.RequireRecvCtx(ctx, t, client2.resps)
testutil.RequireSendCtx(ctx, t, respCall.err, io.EOF)
closeCall = testutil.RequireRecvCtx(ctx, t, client2.close)
testutil.RequireSendCtx(ctx, t, closeCall, nil)
err = testutil.RequireRecvCtx(ctx, t, cw2.Wait())
respCall = testutil.TryReceive(ctx, t, client2.resps)
testutil.RequireSend(ctx, t, respCall.err, io.EOF)
closeCall = testutil.TryReceive(ctx, t, client2.close)
testutil.RequireSend(ctx, t, closeCall, nil)
err = testutil.TryReceive(ctx, t, cw2.Wait())
require.ErrorIs(t, err, io.EOF)
}
@@ -274,10 +274,10 @@ func TestTunnelSrcCoordController_RemoveDestination_Error(t *testing.T) {
cws <- uut.New(client1)
}()
for range 3 {
call := testutil.RequireRecvCtx(ctx, t, client1.reqs)
testutil.RequireSendCtx(ctx, t, call.err, nil)
call := testutil.TryReceive(ctx, t, client1.reqs)
testutil.RequireSend(ctx, t, call.err, nil)
}
cw1 := testutil.RequireRecvCtx(ctx, t, cws)
cw1 := testutil.TryReceive(ctx, t, cws)
// WHEN: we remove all destinations
removeDone := make(chan struct{})
@@ -290,22 +290,22 @@ func TestTunnelSrcCoordController_RemoveDestination_Error(t *testing.T) {
// WHEN: first RemoveTunnel call fails
theErr := xerrors.New("a bad thing happened")
call := testutil.RequireRecvCtx(ctx, t, client1.reqs)
call := testutil.TryReceive(ctx, t, client1.reqs)
require.Equal(t, dest1[:], call.req.GetRemoveTunnel().GetId())
testutil.RequireSendCtx(ctx, t, call.err, theErr)
testutil.RequireSend(ctx, t, call.err, theErr)
// THEN: we disconnect and do not send remaining RemoveTunnel messages
closeCall := testutil.RequireRecvCtx(ctx, t, client1.close)
testutil.RequireSendCtx(ctx, t, closeCall, nil)
_ = testutil.RequireRecvCtx(ctx, t, removeDone)
closeCall := testutil.TryReceive(ctx, t, client1.close)
testutil.RequireSend(ctx, t, closeCall, nil)
_ = testutil.TryReceive(ctx, t, removeDone)
// shut down
respCall := testutil.RequireRecvCtx(ctx, t, client1.resps)
testutil.RequireSendCtx(ctx, t, respCall.err, io.EOF)
respCall := testutil.TryReceive(ctx, t, client1.resps)
testutil.RequireSend(ctx, t, respCall.err, io.EOF)
// triggers second close call
closeCall = testutil.RequireRecvCtx(ctx, t, client1.close)
testutil.RequireSendCtx(ctx, t, closeCall, nil)
err := testutil.RequireRecvCtx(ctx, t, cw1.Wait())
closeCall = testutil.TryReceive(ctx, t, client1.close)
testutil.RequireSend(ctx, t, closeCall, nil)
err := testutil.TryReceive(ctx, t, cw1.Wait())
require.ErrorIs(t, err, theErr)
}
@@ -331,10 +331,10 @@ func TestTunnelSrcCoordController_Sync(t *testing.T) {
cws <- uut.New(client1)
}()
for range 2 {
call := testutil.RequireRecvCtx(ctx, t, client1.reqs)
testutil.RequireSendCtx(ctx, t, call.err, nil)
call := testutil.TryReceive(ctx, t, client1.reqs)
testutil.RequireSend(ctx, t, call.err, nil)
}
cw1 := testutil.RequireRecvCtx(ctx, t, cws)
cw1 := testutil.TryReceive(ctx, t, cws)
// WHEN: we sync dest2 & dest3
syncDone := make(chan struct{})
@@ -344,23 +344,23 @@ func TestTunnelSrcCoordController_Sync(t *testing.T) {
}()
// THEN: we get an add for dest3 and remove for dest1
call := testutil.RequireRecvCtx(ctx, t, client1.reqs)
call := testutil.TryReceive(ctx, t, client1.reqs)
require.Equal(t, dest3[:], call.req.GetAddTunnel().GetId())
testutil.RequireSendCtx(ctx, t, call.err, nil)
call = testutil.RequireRecvCtx(ctx, t, client1.reqs)
testutil.RequireSend(ctx, t, call.err, nil)
call = testutil.TryReceive(ctx, t, client1.reqs)
require.Equal(t, dest1[:], call.req.GetRemoveTunnel().GetId())
testutil.RequireSendCtx(ctx, t, call.err, nil)
testutil.RequireSend(ctx, t, call.err, nil)
testutil.RequireRecvCtx(ctx, t, syncDone)
testutil.TryReceive(ctx, t, syncDone)
// dest3 should be added to coordinatee
require.Contains(t, fConn.tunnelDestinations, dest3)
// shut down
respCall := testutil.RequireRecvCtx(ctx, t, client1.resps)
testutil.RequireSendCtx(ctx, t, respCall.err, io.EOF)
closeCall := testutil.RequireRecvCtx(ctx, t, client1.close)
testutil.RequireSendCtx(ctx, t, closeCall, nil)
err := testutil.RequireRecvCtx(ctx, t, cw1.Wait())
respCall := testutil.TryReceive(ctx, t, client1.resps)
testutil.RequireSend(ctx, t, respCall.err, io.EOF)
closeCall := testutil.TryReceive(ctx, t, client1.close)
testutil.RequireSend(ctx, t, closeCall, nil)
err := testutil.TryReceive(ctx, t, cw1.Wait())
require.ErrorIs(t, err, io.EOF)
}
@@ -384,24 +384,24 @@ func TestTunnelSrcCoordController_AddDestination_Error(t *testing.T) {
uut.AddDestination(dest1)
}()
theErr := xerrors.New("a bad thing happened")
call := testutil.RequireRecvCtx(ctx, t, client1.reqs)
testutil.RequireSendCtx(ctx, t, call.err, theErr)
call := testutil.TryReceive(ctx, t, client1.reqs)
testutil.RequireSend(ctx, t, call.err, theErr)
// THEN: Client is closed and exits
closeCall := testutil.RequireRecvCtx(ctx, t, client1.close)
testutil.RequireSendCtx(ctx, t, closeCall, nil)
closeCall := testutil.TryReceive(ctx, t, client1.close)
testutil.RequireSend(ctx, t, closeCall, nil)
// close the resps, since the client has closed
resp := testutil.RequireRecvCtx(ctx, t, client1.resps)
testutil.RequireSendCtx(ctx, t, resp.err, net.ErrClosed)
resp := testutil.TryReceive(ctx, t, client1.resps)
testutil.RequireSend(ctx, t, resp.err, net.ErrClosed)
// this triggers a second Close() call on the client
closeCall = testutil.RequireRecvCtx(ctx, t, client1.close)
testutil.RequireSendCtx(ctx, t, closeCall, nil)
closeCall = testutil.TryReceive(ctx, t, client1.close)
testutil.RequireSend(ctx, t, closeCall, nil)
err := testutil.RequireRecvCtx(ctx, t, cw1.Wait())
err := testutil.TryReceive(ctx, t, cw1.Wait())
require.ErrorIs(t, err, theErr)
_ = testutil.RequireRecvCtx(ctx, t, addDone)
_ = testutil.TryReceive(ctx, t, addDone)
}
func TestAgentCoordinationController_SendsReadyForHandshake(t *testing.T) {
@@ -457,7 +457,7 @@ func TestAgentCoordinationController_SendsReadyForHandshake(t *testing.T) {
require.NoError(t, err)
dk, err := key.NewDisco().Public().MarshalText()
require.NoError(t, err)
testutil.RequireSendCtx(ctx, t, resps, &proto.CoordinateResponse{
testutil.RequireSend(ctx, t, resps, &proto.CoordinateResponse{
PeerUpdates: []*proto.CoordinateResponse_PeerUpdate{{
Id: clientID[:],
Kind: proto.CoordinateResponse_PeerUpdate_NODE,
@@ -469,19 +469,19 @@ func TestAgentCoordinationController_SendsReadyForHandshake(t *testing.T) {
}},
})
rfh := testutil.RequireRecvCtx(ctx, t, reqs)
rfh := testutil.TryReceive(ctx, t, reqs)
require.NotNil(t, rfh.ReadyForHandshake)
require.Len(t, rfh.ReadyForHandshake, 1)
require.Equal(t, clientID[:], rfh.ReadyForHandshake[0].Id)
go uut.Close(ctx)
dis := testutil.RequireRecvCtx(ctx, t, reqs)
dis := testutil.TryReceive(ctx, t, reqs)
require.NotNil(t, dis)
require.NotNil(t, dis.Disconnect)
close(resps)
// Recv loop should be terminated by the server hanging up after Disconnect
err = testutil.RequireRecvCtx(ctx, t, uut.Wait())
err = testutil.TryReceive(ctx, t, uut.Wait())
require.ErrorIs(t, err, io.EOF)
}
@@ -493,14 +493,14 @@ func coordinationTest(
agentID uuid.UUID,
) {
// It should add the tunnel, since we configured as a client
req := testutil.RequireRecvCtx(ctx, t, reqs)
req := testutil.TryReceive(ctx, t, reqs)
require.Equal(t, agentID[:], req.GetAddTunnel().GetId())
// when we call the callback, it should send a node update
require.NotNil(t, fConn.callback)
fConn.callback(&tailnet.Node{PreferredDERP: 1})
req = testutil.RequireRecvCtx(ctx, t, reqs)
req = testutil.TryReceive(ctx, t, reqs)
require.Equal(t, int32(1), req.GetUpdateSelf().GetNode().GetPreferredDerp())
// When we send a peer update, it should update the coordinatee
@@ -519,7 +519,7 @@ func coordinationTest(
},
},
}
testutil.RequireSendCtx(ctx, t, resps, &proto.CoordinateResponse{PeerUpdates: updates})
testutil.RequireSend(ctx, t, resps, &proto.CoordinateResponse{PeerUpdates: updates})
require.Eventually(t, func() bool {
fConn.Lock()
defer fConn.Unlock()
@@ -534,11 +534,11 @@ func coordinationTest(
}()
// When we close, it should gracefully disconnect
req = testutil.RequireRecvCtx(ctx, t, reqs)
req = testutil.TryReceive(ctx, t, reqs)
require.NotNil(t, req.Disconnect)
close(resps)
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
require.NoError(t, err)
// It should set all peers lost on the coordinatee
@@ -593,12 +593,12 @@ func TestNewBasicDERPController_Mainline(t *testing.T) {
c := uut.New(fc)
ctx := testutil.Context(t, testutil.WaitShort)
expectDM := &tailcfg.DERPMap{}
testutil.RequireSendCtx(ctx, t, fc.ch, expectDM)
gotDM := testutil.RequireRecvCtx(ctx, t, fs)
testutil.RequireSend(ctx, t, fc.ch, expectDM)
gotDM := testutil.TryReceive(ctx, t, fs)
require.Equal(t, expectDM, gotDM)
err := c.Close(ctx)
require.NoError(t, err)
err = testutil.RequireRecvCtx(ctx, t, c.Wait())
err = testutil.TryReceive(ctx, t, c.Wait())
require.ErrorIs(t, err, io.EOF)
// ensure Close is idempotent
err = c.Close(ctx)
@@ -617,7 +617,7 @@ func TestNewBasicDERPController_RecvErr(t *testing.T) {
}
c := uut.New(fc)
ctx := testutil.Context(t, testutil.WaitShort)
err := testutil.RequireRecvCtx(ctx, t, c.Wait())
err := testutil.TryReceive(ctx, t, c.Wait())
require.ErrorIs(t, err, expectedErr)
// ensure Close is idempotent
err = c.Close(ctx)
@@ -668,12 +668,12 @@ func TestBasicTelemetryController_Success(t *testing.T) {
})
}()
call := testutil.RequireRecvCtx(ctx, t, ft.calls)
call := testutil.TryReceive(ctx, t, ft.calls)
require.Len(t, call.req.GetEvents(), 1)
require.Equal(t, call.req.GetEvents()[0].GetId(), []byte("test event"))
testutil.RequireSendCtx(ctx, t, call.errCh, nil)
testutil.RequireRecvCtx(ctx, t, sendDone)
testutil.RequireSend(ctx, t, call.errCh, nil)
testutil.TryReceive(ctx, t, sendDone)
}
func TestBasicTelemetryController_Unimplemented(t *testing.T) {
@@ -695,9 +695,9 @@ func TestBasicTelemetryController_Unimplemented(t *testing.T) {
uut.SendTelemetryEvent(&proto.TelemetryEvent{})
}()
call := testutil.RequireRecvCtx(ctx, t, ft.calls)
testutil.RequireSendCtx(ctx, t, call.errCh, telemetryError)
testutil.RequireRecvCtx(ctx, t, sendDone)
call := testutil.TryReceive(ctx, t, ft.calls)
testutil.RequireSend(ctx, t, call.errCh, telemetryError)
testutil.TryReceive(ctx, t, sendDone)
sendDone = make(chan struct{})
go func() {
@@ -706,12 +706,12 @@ func TestBasicTelemetryController_Unimplemented(t *testing.T) {
}()
// we get another call since it wasn't really the Unimplemented error
call = testutil.RequireRecvCtx(ctx, t, ft.calls)
call = testutil.TryReceive(ctx, t, ft.calls)
// for real this time
telemetryError = errUnimplemented
testutil.RequireSendCtx(ctx, t, call.errCh, telemetryError)
testutil.RequireRecvCtx(ctx, t, sendDone)
testutil.RequireSend(ctx, t, call.errCh, telemetryError)
testutil.TryReceive(ctx, t, sendDone)
// now this returns immediately without a call, because unimplemented error disables calling
sendDone = make(chan struct{})
@@ -719,7 +719,7 @@ func TestBasicTelemetryController_Unimplemented(t *testing.T) {
defer close(sendDone)
uut.SendTelemetryEvent(&proto.TelemetryEvent{})
}()
testutil.RequireRecvCtx(ctx, t, sendDone)
testutil.TryReceive(ctx, t, sendDone)
// getting a "new" client resets
uut.New(ft)
@@ -728,9 +728,9 @@ func TestBasicTelemetryController_Unimplemented(t *testing.T) {
defer close(sendDone)
uut.SendTelemetryEvent(&proto.TelemetryEvent{})
}()
call = testutil.RequireRecvCtx(ctx, t, ft.calls)
testutil.RequireSendCtx(ctx, t, call.errCh, nil)
testutil.RequireRecvCtx(ctx, t, sendDone)
call = testutil.TryReceive(ctx, t, ft.calls)
testutil.RequireSend(ctx, t, call.errCh, nil)
testutil.TryReceive(ctx, t, sendDone)
}
func TestBasicTelemetryController_NotRecognised(t *testing.T) {
@@ -747,20 +747,20 @@ func TestBasicTelemetryController_NotRecognised(t *testing.T) {
uut.SendTelemetryEvent(&proto.TelemetryEvent{})
}()
// returning generic protocol error doesn't trigger unknown rpc logic
call := testutil.RequireRecvCtx(ctx, t, ft.calls)
testutil.RequireSendCtx(ctx, t, call.errCh, drpc.ProtocolError.New("Protocol Error"))
testutil.RequireRecvCtx(ctx, t, sendDone)
call := testutil.TryReceive(ctx, t, ft.calls)
testutil.RequireSend(ctx, t, call.errCh, drpc.ProtocolError.New("Protocol Error"))
testutil.TryReceive(ctx, t, sendDone)
sendDone = make(chan struct{})
go func() {
defer close(sendDone)
uut.SendTelemetryEvent(&proto.TelemetryEvent{})
}()
call = testutil.RequireRecvCtx(ctx, t, ft.calls)
call = testutil.TryReceive(ctx, t, ft.calls)
// return the expected protocol error this time
testutil.RequireSendCtx(ctx, t, call.errCh,
testutil.RequireSend(ctx, t, call.errCh,
drpc.ProtocolError.New("unknown rpc: /coder.tailnet.v2.Tailnet/PostTelemetry"))
testutil.RequireRecvCtx(ctx, t, sendDone)
testutil.TryReceive(ctx, t, sendDone)
// now this returns immediately without a call, because unimplemented error disables calling
sendDone = make(chan struct{})
@@ -768,7 +768,7 @@ func TestBasicTelemetryController_NotRecognised(t *testing.T) {
defer close(sendDone)
uut.SendTelemetryEvent(&proto.TelemetryEvent{})
}()
testutil.RequireRecvCtx(ctx, t, sendDone)
testutil.TryReceive(ctx, t, sendDone)
}
type fakeTelemetryClient struct {
@@ -822,8 +822,8 @@ func TestBasicResumeTokenController_Mainline(t *testing.T) {
go func() {
cwCh <- uut.New(fr)
}()
call := testutil.RequireRecvCtx(ctx, t, fr.calls)
testutil.RequireSendCtx(ctx, t, call.resp, &proto.RefreshResumeTokenResponse{
call := testutil.TryReceive(ctx, t, fr.calls)
testutil.RequireSend(ctx, t, call.resp, &proto.RefreshResumeTokenResponse{
Token: "test token 1",
RefreshIn: durationpb.New(100 * time.Second),
ExpiresAt: timestamppb.New(mClock.Now().Add(200 * time.Second)),
@@ -832,11 +832,11 @@ func TestBasicResumeTokenController_Mainline(t *testing.T) {
token, ok := uut.Token()
require.True(t, ok)
require.Equal(t, "test token 1", token)
cw := testutil.RequireRecvCtx(ctx, t, cwCh)
cw := testutil.TryReceive(ctx, t, cwCh)
w := mClock.Advance(100 * time.Second)
call = testutil.RequireRecvCtx(ctx, t, fr.calls)
testutil.RequireSendCtx(ctx, t, call.resp, &proto.RefreshResumeTokenResponse{
call = testutil.TryReceive(ctx, t, fr.calls)
testutil.RequireSend(ctx, t, call.resp, &proto.RefreshResumeTokenResponse{
Token: "test token 2",
RefreshIn: durationpb.New(50 * time.Second),
ExpiresAt: timestamppb.New(mClock.Now().Add(200 * time.Second)),
@@ -851,7 +851,7 @@ func TestBasicResumeTokenController_Mainline(t *testing.T) {
err := cw.Close(ctx)
require.NoError(t, err)
err = testutil.RequireRecvCtx(ctx, t, cw.Wait())
err = testutil.TryReceive(ctx, t, cw.Wait())
require.NoError(t, err)
token, ok = uut.Token()
@@ -880,24 +880,24 @@ func TestBasicResumeTokenController_NewWhileRefreshing(t *testing.T) {
go func() {
cwCh1 <- uut.New(fr1)
}()
call1 := testutil.RequireRecvCtx(ctx, t, fr1.calls)
call1 := testutil.TryReceive(ctx, t, fr1.calls)
fr2 := newFakeResumeTokenClient(ctx)
cwCh2 := make(chan tailnet.CloserWaiter, 1)
go func() {
cwCh2 <- uut.New(fr2)
}()
call2 := testutil.RequireRecvCtx(ctx, t, fr2.calls)
call2 := testutil.TryReceive(ctx, t, fr2.calls)
testutil.RequireSendCtx(ctx, t, call2.resp, &proto.RefreshResumeTokenResponse{
testutil.RequireSend(ctx, t, call2.resp, &proto.RefreshResumeTokenResponse{
Token: "test token 2.0",
RefreshIn: durationpb.New(102 * time.Second),
ExpiresAt: timestamppb.New(mClock.Now().Add(200 * time.Second)),
})
cw2 := testutil.RequireRecvCtx(ctx, t, cwCh2) // this ensures Close was called on 1
cw2 := testutil.TryReceive(ctx, t, cwCh2) // this ensures Close was called on 1
testutil.RequireSendCtx(ctx, t, call1.resp, &proto.RefreshResumeTokenResponse{
testutil.RequireSend(ctx, t, call1.resp, &proto.RefreshResumeTokenResponse{
Token: "test token 1",
RefreshIn: durationpb.New(101 * time.Second),
ExpiresAt: timestamppb.New(mClock.Now().Add(200 * time.Second)),
@@ -910,13 +910,13 @@ func TestBasicResumeTokenController_NewWhileRefreshing(t *testing.T) {
require.Equal(t, "test token 2.0", token)
// refresher 1 should already be closed.
cw1 := testutil.RequireRecvCtx(ctx, t, cwCh1)
err := testutil.RequireRecvCtx(ctx, t, cw1.Wait())
cw1 := testutil.TryReceive(ctx, t, cwCh1)
err := testutil.TryReceive(ctx, t, cw1.Wait())
require.NoError(t, err)
w := mClock.Advance(102 * time.Second)
call := testutil.RequireRecvCtx(ctx, t, fr2.calls)
testutil.RequireSendCtx(ctx, t, call.resp, &proto.RefreshResumeTokenResponse{
call := testutil.TryReceive(ctx, t, fr2.calls)
testutil.RequireSend(ctx, t, call.resp, &proto.RefreshResumeTokenResponse{
Token: "test token 2.1",
RefreshIn: durationpb.New(50 * time.Second),
ExpiresAt: timestamppb.New(mClock.Now().Add(200 * time.Second)),
@@ -931,7 +931,7 @@ func TestBasicResumeTokenController_NewWhileRefreshing(t *testing.T) {
err = cw2.Close(ctx)
require.NoError(t, err)
err = testutil.RequireRecvCtx(ctx, t, cw2.Wait())
err = testutil.TryReceive(ctx, t, cw2.Wait())
require.NoError(t, err)
}
@@ -948,9 +948,9 @@ func TestBasicResumeTokenController_Unimplemented(t *testing.T) {
fr := newFakeResumeTokenClient(ctx)
cw := uut.New(fr)
call := testutil.RequireRecvCtx(ctx, t, fr.calls)
testutil.RequireSendCtx(ctx, t, call.errCh, errUnimplemented)
err := testutil.RequireRecvCtx(ctx, t, cw.Wait())
call := testutil.TryReceive(ctx, t, fr.calls)
testutil.RequireSend(ctx, t, call.errCh, errUnimplemented)
err := testutil.TryReceive(ctx, t, cw.Wait())
require.NoError(t, err)
_, ok = uut.Token()
require.False(t, ok)
@@ -1044,35 +1044,35 @@ func TestController_Disconnects(t *testing.T) {
uut.DERPCtrl = tailnet.NewBasicDERPController(logger.Named("derp_ctrl"), fConn)
uut.Run(ctx)
call := testutil.RequireRecvCtx(testCtx, t, fCoord.CoordinateCalls)
call := testutil.TryReceive(testCtx, t, fCoord.CoordinateCalls)
// simulate a problem with DERPMaps by sending nil
testutil.RequireSendCtx(testCtx, t, derpMapCh, nil)
testutil.RequireSend(testCtx, t, derpMapCh, nil)
// this should cause the coordinate call to hang up WITHOUT disconnecting
reqNil := testutil.RequireRecvCtx(testCtx, t, call.Reqs)
reqNil := testutil.TryReceive(testCtx, t, call.Reqs)
require.Nil(t, reqNil)
// and mark all peers lost
_ = testutil.RequireRecvCtx(testCtx, t, peersLost)
_ = testutil.TryReceive(testCtx, t, peersLost)
// ...and then reconnect
call = testutil.RequireRecvCtx(testCtx, t, fCoord.CoordinateCalls)
call = testutil.TryReceive(testCtx, t, fCoord.CoordinateCalls)
// close the coordination call, which should cause a 2nd reconnection
close(call.Resps)
_ = testutil.RequireRecvCtx(testCtx, t, peersLost)
call = testutil.RequireRecvCtx(testCtx, t, fCoord.CoordinateCalls)
_ = testutil.TryReceive(testCtx, t, peersLost)
call = testutil.TryReceive(testCtx, t, fCoord.CoordinateCalls)
// canceling the context should trigger the disconnect message
cancel()
reqDisc := testutil.RequireRecvCtx(testCtx, t, call.Reqs)
reqDisc := testutil.TryReceive(testCtx, t, call.Reqs)
require.NotNil(t, reqDisc)
require.NotNil(t, reqDisc.Disconnect)
close(call.Resps)
_ = testutil.RequireRecvCtx(testCtx, t, peersLost)
_ = testutil.RequireRecvCtx(testCtx, t, uut.Closed())
_ = testutil.TryReceive(testCtx, t, peersLost)
_ = testutil.TryReceive(testCtx, t, uut.Closed())
}
func TestController_TelemetrySuccess(t *testing.T) {
@@ -1124,14 +1124,14 @@ func TestController_TelemetrySuccess(t *testing.T) {
uut.Run(ctx)
// Coordinate calls happen _after_ telemetry is connected up, so we use this
// to ensure telemetry is connected before sending our event
cc := testutil.RequireRecvCtx(ctx, t, fCoord.CoordinateCalls)
cc := testutil.TryReceive(ctx, t, fCoord.CoordinateCalls)
defer close(cc.Resps)
tel.SendTelemetryEvent(&proto.TelemetryEvent{
Id: []byte("test event"),
})
testEvents := testutil.RequireRecvCtx(ctx, t, eventCh)
testEvents := testutil.TryReceive(ctx, t, eventCh)
require.Len(t, testEvents, 1)
require.Equal(t, []byte("test event"), testEvents[0].Id)
@@ -1157,27 +1157,27 @@ func TestController_WorkspaceUpdates(t *testing.T) {
uut.Run(ctx)
// it should dial and pass the client to the controller
call := testutil.RequireRecvCtx(testCtx, t, fCtrl.calls)
call := testutil.TryReceive(testCtx, t, fCtrl.calls)
require.Equal(t, fClient, call.client)
fCW := newFakeCloserWaiter()
testutil.RequireSendCtx[tailnet.CloserWaiter](testCtx, t, call.resp, fCW)
testutil.RequireSend[tailnet.CloserWaiter](testCtx, t, call.resp, fCW)
// if the CloserWaiter exits...
testutil.RequireSendCtx(testCtx, t, fCW.errCh, theError)
testutil.RequireSend(testCtx, t, fCW.errCh, theError)
// it should close, redial and reconnect
cCall := testutil.RequireRecvCtx(testCtx, t, fClient.close)
testutil.RequireSendCtx(testCtx, t, cCall, nil)
cCall := testutil.TryReceive(testCtx, t, fClient.close)
testutil.RequireSend(testCtx, t, cCall, nil)
call = testutil.RequireRecvCtx(testCtx, t, fCtrl.calls)
call = testutil.TryReceive(testCtx, t, fCtrl.calls)
require.Equal(t, fClient, call.client)
fCW = newFakeCloserWaiter()
testutil.RequireSendCtx[tailnet.CloserWaiter](testCtx, t, call.resp, fCW)
testutil.RequireSend[tailnet.CloserWaiter](testCtx, t, call.resp, fCW)
// canceling the context should close the client
cancel()
cCall = testutil.RequireRecvCtx(testCtx, t, fClient.close)
testutil.RequireSendCtx(testCtx, t, cCall, nil)
cCall = testutil.TryReceive(testCtx, t, fClient.close)
testutil.RequireSend(testCtx, t, cCall, nil)
}
type fakeTailnetConn struct {
@@ -1492,12 +1492,12 @@ func setupConnectedAllWorkspaceUpdatesController(
coordCW := tsc.New(coordC)
t.Cleanup(func() {
// hang up coord client
coordRecv := testutil.RequireRecvCtx(ctx, t, coordC.resps)
testutil.RequireSendCtx(ctx, t, coordRecv.err, io.EOF)
coordRecv := testutil.TryReceive(ctx, t, coordC.resps)
testutil.RequireSend(ctx, t, coordRecv.err, io.EOF)
// sends close on client
cCall := testutil.RequireRecvCtx(ctx, t, coordC.close)
testutil.RequireSendCtx(ctx, t, cCall, nil)
err := testutil.RequireRecvCtx(ctx, t, coordCW.Wait())
cCall := testutil.TryReceive(ctx, t, coordC.close)
testutil.RequireSend(ctx, t, cCall, nil)
err := testutil.TryReceive(ctx, t, coordCW.Wait())
require.ErrorIs(t, err, io.EOF)
})
@@ -1506,9 +1506,9 @@ func setupConnectedAllWorkspaceUpdatesController(
updateCW := uut.New(updateC)
t.Cleanup(func() {
// hang up WorkspaceUpdates client
upRecvCall := testutil.RequireRecvCtx(ctx, t, updateC.recv)
testutil.RequireSendCtx(ctx, t, upRecvCall.err, io.EOF)
err := testutil.RequireRecvCtx(ctx, t, updateCW.Wait())
upRecvCall := testutil.TryReceive(ctx, t, updateC.recv)
testutil.RequireSend(ctx, t, upRecvCall.err, io.EOF)
err := testutil.TryReceive(ctx, t, updateCW.Wait())
require.ErrorIs(t, err, io.EOF)
})
return coordC, updateC, uut
@@ -1544,15 +1544,15 @@ func TestTunnelAllWorkspaceUpdatesController_Initial(t *testing.T) {
},
}
upRecvCall := testutil.RequireRecvCtx(ctx, t, updateC.recv)
testutil.RequireSendCtx(ctx, t, upRecvCall.resp, initUp)
upRecvCall := testutil.TryReceive(ctx, t, updateC.recv)
testutil.RequireSend(ctx, t, upRecvCall.resp, initUp)
// This should trigger AddTunnel for each agent
var adds []uuid.UUID
for range 3 {
coordCall := testutil.RequireRecvCtx(ctx, t, coordC.reqs)
coordCall := testutil.TryReceive(ctx, t, coordC.reqs)
adds = append(adds, uuid.Must(uuid.FromBytes(coordCall.req.GetAddTunnel().GetId())))
testutil.RequireSendCtx(ctx, t, coordCall.err, nil)
testutil.RequireSend(ctx, t, coordCall.err, nil)
}
require.Contains(t, adds, w1a1ID)
require.Contains(t, adds, w2a1ID)
@@ -1576,9 +1576,9 @@ func TestTunnelAllWorkspaceUpdatesController_Initial(t *testing.T) {
"w1.mctest.": {ws1a1IP},
expectedCoderConnectFQDN: {tsaddr.CoderServiceIPv6()},
}
dnsCall := testutil.RequireRecvCtx(ctx, t, fDNS.calls)
dnsCall := testutil.TryReceive(ctx, t, fDNS.calls)
require.Equal(t, expectedDNS, dnsCall.hosts)
testutil.RequireSendCtx(ctx, t, dnsCall.err, nil)
testutil.RequireSend(ctx, t, dnsCall.err, nil)
currentState := tailnet.WorkspaceUpdate{
UpsertedWorkspaces: []*tailnet.Workspace{
@@ -1614,7 +1614,7 @@ func TestTunnelAllWorkspaceUpdatesController_Initial(t *testing.T) {
}
// And the callback
cbUpdate := testutil.RequireRecvCtx(ctx, t, fUH.ch)
cbUpdate := testutil.TryReceive(ctx, t, fUH.ch)
require.Equal(t, currentState, cbUpdate)
// Current recvState should match
@@ -1656,13 +1656,13 @@ func TestTunnelAllWorkspaceUpdatesController_DeleteAgent(t *testing.T) {
},
}
upRecvCall := testutil.RequireRecvCtx(ctx, t, updateC.recv)
testutil.RequireSendCtx(ctx, t, upRecvCall.resp, initUp)
upRecvCall := testutil.TryReceive(ctx, t, updateC.recv)
testutil.RequireSend(ctx, t, upRecvCall.resp, initUp)
// Add for w1a1
coordCall := testutil.RequireRecvCtx(ctx, t, coordC.reqs)
coordCall := testutil.TryReceive(ctx, t, coordC.reqs)
require.Equal(t, w1a1ID[:], coordCall.req.GetAddTunnel().GetId())
testutil.RequireSendCtx(ctx, t, coordCall.err, nil)
testutil.RequireSend(ctx, t, coordCall.err, nil)
expectedCoderConnectFQDN, err := dnsname.ToFQDN(
fmt.Sprintf(tailnet.IsCoderConnectEnabledFmtString, tailnet.CoderDNSSuffix))
@@ -1675,9 +1675,9 @@ func TestTunnelAllWorkspaceUpdatesController_DeleteAgent(t *testing.T) {
"w1.coder.": {ws1a1IP},
expectedCoderConnectFQDN: {tsaddr.CoderServiceIPv6()},
}
dnsCall := testutil.RequireRecvCtx(ctx, t, fDNS.calls)
dnsCall := testutil.TryReceive(ctx, t, fDNS.calls)
require.Equal(t, expectedDNS, dnsCall.hosts)
testutil.RequireSendCtx(ctx, t, dnsCall.err, nil)
testutil.RequireSend(ctx, t, dnsCall.err, nil)
initRecvUp := tailnet.WorkspaceUpdate{
UpsertedWorkspaces: []*tailnet.Workspace{
@@ -1694,7 +1694,7 @@ func TestTunnelAllWorkspaceUpdatesController_DeleteAgent(t *testing.T) {
DeletedAgents: []*tailnet.Agent{},
}
cbUpdate := testutil.RequireRecvCtx(ctx, t, fUH.ch)
cbUpdate := testutil.TryReceive(ctx, t, fUH.ch)
require.Equal(t, initRecvUp, cbUpdate)
// Current state should match initial
@@ -1711,18 +1711,18 @@ func TestTunnelAllWorkspaceUpdatesController_DeleteAgent(t *testing.T) {
{Id: w1a1ID[:], WorkspaceId: w1ID[:]},
},
}
upRecvCall = testutil.RequireRecvCtx(ctx, t, updateC.recv)
testutil.RequireSendCtx(ctx, t, upRecvCall.resp, agentUpdate)
upRecvCall = testutil.TryReceive(ctx, t, updateC.recv)
testutil.RequireSend(ctx, t, upRecvCall.resp, agentUpdate)
// Add for w1a2
coordCall = testutil.RequireRecvCtx(ctx, t, coordC.reqs)
coordCall = testutil.TryReceive(ctx, t, coordC.reqs)
require.Equal(t, w1a2ID[:], coordCall.req.GetAddTunnel().GetId())
testutil.RequireSendCtx(ctx, t, coordCall.err, nil)
testutil.RequireSend(ctx, t, coordCall.err, nil)
// Remove for w1a1
coordCall = testutil.RequireRecvCtx(ctx, t, coordC.reqs)
coordCall = testutil.TryReceive(ctx, t, coordC.reqs)
require.Equal(t, w1a1ID[:], coordCall.req.GetRemoveTunnel().GetId())
testutil.RequireSendCtx(ctx, t, coordCall.err, nil)
testutil.RequireSend(ctx, t, coordCall.err, nil)
// DNS contains only w1a2
expectedDNS = map[dnsname.FQDN][]netip.Addr{
@@ -1731,11 +1731,11 @@ func TestTunnelAllWorkspaceUpdatesController_DeleteAgent(t *testing.T) {
"w1.coder.": {ws1a2IP},
expectedCoderConnectFQDN: {tsaddr.CoderServiceIPv6()},
}
dnsCall = testutil.RequireRecvCtx(ctx, t, fDNS.calls)
dnsCall = testutil.TryReceive(ctx, t, fDNS.calls)
require.Equal(t, expectedDNS, dnsCall.hosts)
testutil.RequireSendCtx(ctx, t, dnsCall.err, nil)
testutil.RequireSend(ctx, t, dnsCall.err, nil)
cbUpdate = testutil.RequireRecvCtx(ctx, t, fUH.ch)
cbUpdate = testutil.TryReceive(ctx, t, fUH.ch)
sndRecvUpdate := tailnet.WorkspaceUpdate{
UpsertedWorkspaces: []*tailnet.Workspace{},
UpsertedAgents: []*tailnet.Agent{
@@ -1804,8 +1804,8 @@ func TestTunnelAllWorkspaceUpdatesController_DNSError(t *testing.T) {
{Id: w1a1ID[:], Name: "w1a1", WorkspaceId: w1ID[:]},
},
}
upRecvCall := testutil.RequireRecvCtx(ctx, t, updateC.recv)
testutil.RequireSendCtx(ctx, t, upRecvCall.resp, initUp)
upRecvCall := testutil.TryReceive(ctx, t, updateC.recv)
testutil.RequireSend(ctx, t, upRecvCall.resp, initUp)
expectedCoderConnectFQDN, err := dnsname.ToFQDN(
fmt.Sprintf(tailnet.IsCoderConnectEnabledFmtString, tailnet.CoderDNSSuffix))
@@ -1818,16 +1818,16 @@ func TestTunnelAllWorkspaceUpdatesController_DNSError(t *testing.T) {
"w1.coder.": {ws1a1IP},
expectedCoderConnectFQDN: {tsaddr.CoderServiceIPv6()},
}
dnsCall := testutil.RequireRecvCtx(ctx, t, fDNS.calls)
dnsCall := testutil.TryReceive(ctx, t, fDNS.calls)
require.Equal(t, expectedDNS, dnsCall.hosts)
testutil.RequireSendCtx(ctx, t, dnsCall.err, dnsError)
testutil.RequireSend(ctx, t, dnsCall.err, dnsError)
// should trigger a close on the client
closeCall := testutil.RequireRecvCtx(ctx, t, updateC.close)
testutil.RequireSendCtx(ctx, t, closeCall, io.EOF)
closeCall := testutil.TryReceive(ctx, t, updateC.close)
testutil.RequireSend(ctx, t, closeCall, io.EOF)
// error should be our initial DNS error
err = testutil.RequireRecvCtx(ctx, t, updateCW.Wait())
err = testutil.TryReceive(ctx, t, updateCW.Wait())
require.ErrorIs(t, err, dnsError)
}
@@ -1927,12 +1927,12 @@ func TestTunnelAllWorkspaceUpdatesController_HandleErrors(t *testing.T) {
updateC := newFakeWorkspaceUpdateClient(ctx, t)
updateCW := uut.New(updateC)
recvCall := testutil.RequireRecvCtx(ctx, t, updateC.recv)
testutil.RequireSendCtx(ctx, t, recvCall.resp, tc.update)
closeCall := testutil.RequireRecvCtx(ctx, t, updateC.close)
testutil.RequireSendCtx(ctx, t, closeCall, nil)
recvCall := testutil.TryReceive(ctx, t, updateC.recv)
testutil.RequireSend(ctx, t, recvCall.resp, tc.update)
closeCall := testutil.TryReceive(ctx, t, updateC.close)
testutil.RequireSend(ctx, t, closeCall, nil)
err := testutil.RequireRecvCtx(ctx, t, updateCW.Wait())
err := testutil.TryReceive(ctx, t, updateCW.Wait())
require.ErrorContains(t, err, tc.errorContains)
})
}
tailnet/coordinator_test.go
+2 -2
View File
@@ -270,8 +270,8 @@ func TestCoordinatorPropogatedPeerContext(t *testing.T) {
reqs, _ := c1.Coordinate(peerCtx, peerID, "peer1", auth)
testutil.RequireSendCtx(ctx, t, reqs, &proto.CoordinateRequest{AddTunnel: &proto.CoordinateRequest_Tunnel{Id: tailnet.UUIDToByteSlice(agentID)}})
_ = testutil.RequireRecvCtx(ctx, t, ch)
testutil.RequireSend(ctx, t, reqs, &proto.CoordinateRequest{AddTunnel: &proto.CoordinateRequest_Tunnel{Id: tailnet.UUIDToByteSlice(agentID)}})
_ = testutil.TryReceive(ctx, t, ch)
// If we don't cancel the context, the coordinator close will wait until the
// peer request loop finishes, which will be after the timeout
peerCtxCancel()
tailnet/node_internal_test.go
+23 -23
View File
@@ -42,7 +42,7 @@ func TestNodeUpdater_setNetInfo_different(t *testing.T) {
DERPLatency: dl,
})
node := testutil.RequireRecvCtx(ctx, t, nodeCh)
node := testutil.TryReceive(ctx, t, nodeCh)
require.Equal(t, nodeKey, node.Key)
require.Equal(t, discoKey, node.DiscoKey)
require.Equal(t, 1, node.PreferredDERP)
@@ -56,7 +56,7 @@ func TestNodeUpdater_setNetInfo_different(t *testing.T) {
})
close(goCh) // allows callback to complete
node = testutil.RequireRecvCtx(ctx, t, nodeCh)
node = testutil.TryReceive(ctx, t, nodeCh)
require.Equal(t, nodeKey, node.Key)
require.Equal(t, discoKey, node.DiscoKey)
require.Equal(t, 2, node.PreferredDERP)
@@ -67,7 +67,7 @@ func TestNodeUpdater_setNetInfo_different(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestNodeUpdater_setNetInfo_same(t *testing.T) {
@@ -108,7 +108,7 @@ func TestNodeUpdater_setNetInfo_same(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestNodeUpdater_setDERPForcedWebsocket_different(t *testing.T) {
@@ -137,7 +137,7 @@ func TestNodeUpdater_setDERPForcedWebsocket_different(t *testing.T) {
uut.setDERPForcedWebsocket(1, "test")
// Then: we receive an update with the reason set
node := testutil.RequireRecvCtx(ctx, t, nodeCh)
node := testutil.TryReceive(ctx, t, nodeCh)
require.Equal(t, nodeKey, node.Key)
require.Equal(t, discoKey, node.DiscoKey)
require.True(t, maps.Equal(map[int]string{1: "test"}, node.DERPForcedWebsocket))
@@ -147,7 +147,7 @@ func TestNodeUpdater_setDERPForcedWebsocket_different(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestNodeUpdater_setDERPForcedWebsocket_same(t *testing.T) {
@@ -185,7 +185,7 @@ func TestNodeUpdater_setDERPForcedWebsocket_same(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestNodeUpdater_setStatus_different(t *testing.T) {
@@ -220,7 +220,7 @@ func TestNodeUpdater_setStatus_different(t *testing.T) {
}, nil)
// Then: we receive an update with the endpoint
node := testutil.RequireRecvCtx(ctx, t, nodeCh)
node := testutil.TryReceive(ctx, t, nodeCh)
require.Equal(t, nodeKey, node.Key)
require.Equal(t, discoKey, node.DiscoKey)
require.Equal(t, []string{"[fe80::1]:5678"}, node.Endpoints)
@@ -235,7 +235,7 @@ func TestNodeUpdater_setStatus_different(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestNodeUpdater_setStatus_same(t *testing.T) {
@@ -275,7 +275,7 @@ func TestNodeUpdater_setStatus_same(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestNodeUpdater_setStatus_error(t *testing.T) {
@@ -313,7 +313,7 @@ func TestNodeUpdater_setStatus_error(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestNodeUpdater_setStatus_outdated(t *testing.T) {
@@ -355,7 +355,7 @@ func TestNodeUpdater_setStatus_outdated(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestNodeUpdater_setAddresses_different(t *testing.T) {
@@ -385,7 +385,7 @@ func TestNodeUpdater_setAddresses_different(t *testing.T) {
uut.setAddresses(addrs)
// Then: we receive an update with the addresses
node := testutil.RequireRecvCtx(ctx, t, nodeCh)
node := testutil.TryReceive(ctx, t, nodeCh)
require.Equal(t, nodeKey, node.Key)
require.Equal(t, discoKey, node.DiscoKey)
require.Equal(t, addrs, node.Addresses)
@@ -396,7 +396,7 @@ func TestNodeUpdater_setAddresses_different(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestNodeUpdater_setAddresses_same(t *testing.T) {
@@ -435,7 +435,7 @@ func TestNodeUpdater_setAddresses_same(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestNodeUpdater_setCallback(t *testing.T) {
@@ -466,7 +466,7 @@ func TestNodeUpdater_setCallback(t *testing.T) {
})
// Then: we get a node update
node := testutil.RequireRecvCtx(ctx, t, nodeCh)
node := testutil.TryReceive(ctx, t, nodeCh)
require.Equal(t, nodeKey, node.Key)
require.Equal(t, discoKey, node.DiscoKey)
require.Equal(t, 1, node.PreferredDERP)
@@ -476,7 +476,7 @@ func TestNodeUpdater_setCallback(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestNodeUpdater_setBlockEndpoints_different(t *testing.T) {
@@ -506,7 +506,7 @@ func TestNodeUpdater_setBlockEndpoints_different(t *testing.T) {
uut.setBlockEndpoints(true)
// Then: we receive an update without endpoints
node := testutil.RequireRecvCtx(ctx, t, nodeCh)
node := testutil.TryReceive(ctx, t, nodeCh)
require.Equal(t, nodeKey, node.Key)
require.Equal(t, discoKey, node.DiscoKey)
require.Len(t, node.Endpoints, 0)
@@ -515,7 +515,7 @@ func TestNodeUpdater_setBlockEndpoints_different(t *testing.T) {
uut.setBlockEndpoints(false)
// Then: we receive an update with endpoints
node = testutil.RequireRecvCtx(ctx, t, nodeCh)
node = testutil.TryReceive(ctx, t, nodeCh)
require.Equal(t, nodeKey, node.Key)
require.Equal(t, discoKey, node.DiscoKey)
require.Len(t, node.Endpoints, 1)
@@ -525,7 +525,7 @@ func TestNodeUpdater_setBlockEndpoints_different(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestNodeUpdater_setBlockEndpoints_same(t *testing.T) {
@@ -563,7 +563,7 @@ func TestNodeUpdater_setBlockEndpoints_same(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestNodeUpdater_fillPeerDiagnostics(t *testing.T) {
@@ -611,7 +611,7 @@ func TestNodeUpdater_fillPeerDiagnostics(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
func TestNodeUpdater_fillPeerDiagnostics_noCallback(t *testing.T) {
@@ -651,5 +651,5 @@ func TestNodeUpdater_fillPeerDiagnostics_noCallback(t *testing.T) {
defer close(done)
uut.close()
}()
_ = testutil.RequireRecvCtx(ctx, t, done)
_ = testutil.TryReceive(ctx, t, done)
}
tailnet/service_test.go
+11 -11
View File
@@ -76,14 +76,14 @@ func TestClientService_ServeClient_V2(t *testing.T) {
})
require.NoError(t, err)
call := testutil.RequireRecvCtx(ctx, t, fCoord.CoordinateCalls)
call := testutil.TryReceive(ctx, t, fCoord.CoordinateCalls)
require.NotNil(t, call)
require.Equal(t, call.ID, clientID)
require.Equal(t, call.Name, "client")
require.NoError(t, call.Auth.Authorize(ctx, &proto.CoordinateRequest{
AddTunnel: &proto.CoordinateRequest_Tunnel{Id: agentID[:]},
}))
req := testutil.RequireRecvCtx(ctx, t, call.Reqs)
req := testutil.TryReceive(ctx, t, call.Reqs)
require.Equal(t, int32(11), req.GetUpdateSelf().GetNode().GetPreferredDerp())
call.Resps <- &proto.CoordinateResponse{PeerUpdates: []*proto.CoordinateResponse_PeerUpdate{
@@ -126,7 +126,7 @@ func TestClientService_ServeClient_V2(t *testing.T) {
res, err := client.PostTelemetry(ctx, telemetryReq)
require.NoError(t, err)
require.NotNil(t, res)
gotEvents := testutil.RequireRecvCtx(ctx, t, telemetryEvents)
gotEvents := testutil.TryReceive(ctx, t, telemetryEvents)
require.Len(t, gotEvents, 2)
require.Equal(t, "hi", string(gotEvents[0].Id))
require.Equal(t, "bye", string(gotEvents[1].Id))
@@ -134,7 +134,7 @@ func TestClientService_ServeClient_V2(t *testing.T) {
// RPCs closed; we need to close the Conn to end the session.
err = c.Close()
require.NoError(t, err)
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
require.True(t, xerrors.Is(err, io.EOF) || xerrors.Is(err, io.ErrClosedPipe))
}
@@ -174,7 +174,7 @@ func TestClientService_ServeClient_V1(t *testing.T) {
errCh <- err
}()
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
require.ErrorIs(t, err, tailnet.ErrUnsupportedVersion)
}
@@ -201,7 +201,7 @@ func TestNetworkTelemetryBatcher(t *testing.T) {
// Should overflow and send a batch.
ctx := testutil.Context(t, testutil.WaitShort)
batch := testutil.RequireRecvCtx(ctx, t, events)
batch := testutil.TryReceive(ctx, t, events)
require.Len(t, batch, 3)
require.Equal(t, "1", string(batch[0].Id))
require.Equal(t, "2", string(batch[1].Id))
@@ -209,7 +209,7 @@ func TestNetworkTelemetryBatcher(t *testing.T) {
// Should send any pending events when the ticker fires.
mClock.Advance(time.Millisecond)
batch = testutil.RequireRecvCtx(ctx, t, events)
batch = testutil.TryReceive(ctx, t, events)
require.Len(t, batch, 1)
require.Equal(t, "4", string(batch[0].Id))
@@ -220,7 +220,7 @@ func TestNetworkTelemetryBatcher(t *testing.T) {
})
err := b.Close()
require.NoError(t, err)
batch = testutil.RequireRecvCtx(ctx, t, events)
batch = testutil.TryReceive(ctx, t, events)
require.Len(t, batch, 2)
require.Equal(t, "5", string(batch[0].Id))
require.Equal(t, "6", string(batch[1].Id))
@@ -250,11 +250,11 @@ func TestClientUserCoordinateeAuth(t *testing.T) {
})
require.NoError(t, err)
call := testutil.RequireRecvCtx(ctx, t, fCoord.CoordinateCalls)
call := testutil.TryReceive(ctx, t, fCoord.CoordinateCalls)
require.NotNil(t, call)
require.Equal(t, call.ID, clientID)
require.Equal(t, call.Name, "client")
req := testutil.RequireRecvCtx(ctx, t, call.Reqs)
req := testutil.TryReceive(ctx, t, call.Reqs)
require.Equal(t, int32(11), req.GetUpdateSelf().GetNode().GetPreferredDerp())
// Authorize uses `ClientUserCoordinateeAuth`
@@ -354,7 +354,7 @@ func createUpdateService(t *testing.T, ctx context.Context, clientID uuid.UUID,
t.Cleanup(func() {
err = c.Close()
require.NoError(t, err)
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
require.True(t, xerrors.Is(err, io.EOF) || xerrors.Is(err, io.ErrClosedPipe))
})
return fCoord, client
testutil/chan.go
+57
View File
@@ -0,0 +1,57 @@
package testutil
import (
"context"
"testing"
)
// TryReceive will attempt to receive a value from the chan and return it. If
// the context expires before a value can be received, it will fail the test. If
// the channel is closed, the zero value of the channel type will be returned.
//
// Safety: Must only be called from the Go routine that created `t`.
func TryReceive[A any](ctx context.Context, t testing.TB, c <-chan A) A {
t.Helper()
select {
case <-ctx.Done():
t.Fatal("timeout")
var a A
return a
case a := <-c:
return a
}
}
// RequireReceive will receive a value from the chan and return it. If the
// context expires or the channel is closed before a value can be received,
// it will fail the test.
//
// Safety: Must only be called from the Go routine that created `t`.
func RequireReceive[A any](ctx context.Context, t testing.TB, c <-chan A) A {
t.Helper()
select {
case <-ctx.Done():
t.Fatal("timeout")
var a A
return a
case a, ok := <-c:
if !ok {
t.Fatal("channel closed")
}
return a
}
}
// RequireSend will send the given value over the chan and then return. If
// the context expires before the send succeeds, it will fail the test.
//
// Safety: Must only be called from the Go routine that created `t`.
func RequireSend[A any](ctx context.Context, t testing.TB, c chan<- A, a A) {
t.Helper()
select {
case <-ctx.Done():
t.Fatal("timeout")
case c <- a:
// OK!
}
}
testutil/ctx.go
-34
View File
@@ -11,37 +11,3 @@ func Context(t *testing.T, dur time.Duration) context.Context {
t.Cleanup(cancel)
return ctx
}
func RequireRecvCtx[A any](ctx context.Context, t testing.TB, c <-chan A) (a A) {
t.Helper()
select {
case <-ctx.Done():
t.Fatal("timeout")
return a
case a = <-c:
return a
}
}
// NOTE: no AssertRecvCtx because it'd be bad if we returned a default value on
// the cases it times out.
func RequireSendCtx[A any](ctx context.Context, t testing.TB, c chan<- A, a A) {
t.Helper()
select {
case <-ctx.Done():
t.Fatal("timeout")
case c <- a:
// OK!
}
}
func AssertSendCtx[A any](ctx context.Context, t testing.TB, c chan<- A, a A) {
t.Helper()
select {
case <-ctx.Done():
t.Error("timeout")
case c <- a:
// OK!
}
}
vpn/client_test.go
+7 -7
View File
@@ -143,11 +143,11 @@ func TestClient_WorkspaceUpdates(t *testing.T) {
connErrCh <- err
connCh <- conn
}()
testutil.RequireRecvCtx(ctx, t, user)
testutil.RequireRecvCtx(ctx, t, connInfo)
err = testutil.RequireRecvCtx(ctx, t, connErrCh)
testutil.TryReceive(ctx, t, user)
testutil.TryReceive(ctx, t, connInfo)
err = testutil.TryReceive(ctx, t, connErrCh)
require.NoError(t, err)
conn := testutil.RequireRecvCtx(ctx, t, connCh)
conn := testutil.TryReceive(ctx, t, connCh)
// Send a workspace update
update := &proto.WorkspaceUpdate{
@@ -165,10 +165,10 @@ func TestClient_WorkspaceUpdates(t *testing.T) {
},
},
}
testutil.RequireSendCtx(ctx, t, outUpdateCh, update)
testutil.RequireSend(ctx, t, outUpdateCh, update)
// It'll be received by the update handler
recvUpdate := testutil.RequireRecvCtx(ctx, t, inUpdateCh)
recvUpdate := testutil.TryReceive(ctx, t, inUpdateCh)
require.Len(t, recvUpdate.UpsertedWorkspaces, 1)
require.Equal(t, wsID, recvUpdate.UpsertedWorkspaces[0].ID)
require.Len(t, recvUpdate.UpsertedAgents, 1)
@@ -202,7 +202,7 @@ func TestClient_WorkspaceUpdates(t *testing.T) {
// Close the conn
conn.Close()
err = testutil.RequireRecvCtx(ctx, t, serveErrCh)
err = testutil.TryReceive(ctx, t, serveErrCh)
require.NoError(t, err)
})
}
vpn/speaker_internal_test.go
+17 -17
View File
@@ -58,12 +58,12 @@ func TestSpeaker_RawPeer(t *testing.T) {
_, err = mp.Write([]byte("codervpn manager 1.3,2.1\n"))
require.NoError(t, err)
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
require.NoError(t, err)
tun.start()
// send a message and verify it follows protocol for encoding
testutil.RequireSendCtx(ctx, t, tun.sendCh, &TunnelMessage{
testutil.RequireSend(ctx, t, tun.sendCh, &TunnelMessage{
Msg: &TunnelMessage_Start{
Start: &StartResponse{},
},
@@ -107,7 +107,7 @@ func TestSpeaker_HandshakeRWFailure(t *testing.T) {
tun = s
errCh <- err
}()
err := testutil.RequireRecvCtx(ctx, t, errCh)
err := testutil.TryReceive(ctx, t, errCh)
require.ErrorContains(t, err, "handshake failed")
require.Nil(t, tun)
}
@@ -131,7 +131,7 @@ func TestSpeaker_HandshakeCtxDone(t *testing.T) {
errCh <- err
}()
cancel()
err := testutil.RequireRecvCtx(testCtx, t, errCh)
err := testutil.TryReceive(testCtx, t, errCh)
require.ErrorContains(t, err, "handshake failed")
require.Nil(t, tun)
}
@@ -168,7 +168,7 @@ func TestSpeaker_OversizeHandshake(t *testing.T) {
_, err = mp.Write([]byte(badHandshake))
require.Error(t, err) // other side closes when we write too much
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
require.ErrorContains(t, err, "handshake failed")
require.Nil(t, tun)
}
@@ -216,7 +216,7 @@ func TestSpeaker_HandshakeInvalid(t *testing.T) {
require.NoError(t, err)
require.Equal(t, expectedHandshake, string(b[:n]))
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
require.ErrorContains(t, err, "validate header")
require.Nil(t, tun)
})
@@ -258,7 +258,7 @@ func TestSpeaker_CorruptMessage(t *testing.T) {
_, err = mp.Write([]byte("codervpn manager 1.0\n"))
require.NoError(t, err)
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
require.NoError(t, err)
tun.start()
@@ -290,7 +290,7 @@ func TestSpeaker_unaryRPC_mainline(t *testing.T) {
resp = r
errCh <- err
}()
req := testutil.RequireRecvCtx(ctx, t, tun.requests)
req := testutil.TryReceive(ctx, t, tun.requests)
require.NotEqualValues(t, 0, req.msg.GetRpc().GetMsgId())
require.Equal(t, "https://coder.example.com", req.msg.GetStart().GetCoderUrl())
err := req.sendReply(&TunnelMessage{
@@ -299,7 +299,7 @@ func TestSpeaker_unaryRPC_mainline(t *testing.T) {
},
})
require.NoError(t, err)
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
require.NoError(t, err)
_, ok := resp.Msg.(*TunnelMessage_Start)
require.True(t, ok)
@@ -334,12 +334,12 @@ func TestSpeaker_unaryRPC_canceled(t *testing.T) {
resp = r
errCh <- err
}()
req := testutil.RequireRecvCtx(testCtx, t, tun.requests)
req := testutil.TryReceive(testCtx, t, tun.requests)
require.NotEqualValues(t, 0, req.msg.GetRpc().GetMsgId())
require.Equal(t, "https://coder.example.com", req.msg.GetStart().GetCoderUrl())
cancel()
err := testutil.RequireRecvCtx(testCtx, t, errCh)
err := testutil.TryReceive(testCtx, t, errCh)
require.ErrorIs(t, err, context.Canceled)
require.Nil(t, resp)
@@ -370,7 +370,7 @@ func TestSpeaker_unaryRPC_hung_up(t *testing.T) {
resp = r
errCh <- err
}()
req := testutil.RequireRecvCtx(testCtx, t, tun.requests)
req := testutil.TryReceive(testCtx, t, tun.requests)
require.NotEqualValues(t, 0, req.msg.GetRpc().GetMsgId())
require.Equal(t, "https://coder.example.com", req.msg.GetStart().GetCoderUrl())
@@ -378,7 +378,7 @@ func TestSpeaker_unaryRPC_hung_up(t *testing.T) {
err := tun.Close()
require.NoError(t, err)
// Then: we should get an error on the RPC.
err = testutil.RequireRecvCtx(testCtx, t, errCh)
err = testutil.TryReceive(testCtx, t, errCh)
require.ErrorIs(t, err, io.ErrUnexpectedEOF)
require.Nil(t, resp)
}
@@ -397,7 +397,7 @@ func TestSpeaker_unaryRPC_sendLoop(t *testing.T) {
// When: serdes sendloop is closed
// Send a message from the manager. This closes the manager serdes sendloop, since it will error
// when writing the message to the (closed) pipe.
testutil.RequireSendCtx(ctx, t, mgr.sendCh, &ManagerMessage{
testutil.RequireSend(ctx, t, mgr.sendCh, &ManagerMessage{
Msg: &ManagerMessage_GetPeerUpdate{},
})
@@ -417,7 +417,7 @@ func TestSpeaker_unaryRPC_sendLoop(t *testing.T) {
}()
// Then: we should get an error on the RPC.
err = testutil.RequireRecvCtx(testCtx, t, errCh)
err = testutil.TryReceive(testCtx, t, errCh)
require.ErrorIs(t, err, io.ErrUnexpectedEOF)
require.Nil(t, resp)
}
@@ -448,9 +448,9 @@ func setupSpeakers(t *testing.T) (
mgr = s
errCh <- err
}()
err := testutil.RequireRecvCtx(ctx, t, errCh)
err := testutil.TryReceive(ctx, t, errCh)
require.NoError(t, err)
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
require.NoError(t, err)
tun.start()
mgr.start()
vpn/tunnel_internal_test.go
+23 -23
View File
@@ -114,9 +114,9 @@ func TestTunnel_StartStop(t *testing.T) {
errCh <- err
}()
// Then: `NewConn` is called,
testutil.RequireSendCtx(ctx, t, client.ch, conn)
testutil.RequireSend(ctx, t, client.ch, conn)
// And: a response is received
err := testutil.RequireRecvCtx(ctx, t, errCh)
err := testutil.TryReceive(ctx, t, errCh)
require.NoError(t, err)
_, ok := resp.Msg.(*TunnelMessage_Start)
require.True(t, ok)
@@ -130,9 +130,9 @@ func TestTunnel_StartStop(t *testing.T) {
errCh <- err
}()
// Then: `Close` is called on the connection
testutil.RequireRecvCtx(ctx, t, conn.closed)
testutil.TryReceive(ctx, t, conn.closed)
// And: a Stop response is received
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
require.NoError(t, err)
_, ok = resp.Msg.(*TunnelMessage_Stop)
require.True(t, ok)
@@ -178,8 +178,8 @@ func TestTunnel_PeerUpdate(t *testing.T) {
resp = r
errCh <- err
}()
testutil.RequireSendCtx(ctx, t, client.ch, conn)
err := testutil.RequireRecvCtx(ctx, t, errCh)
testutil.RequireSend(ctx, t, client.ch, conn)
err := testutil.TryReceive(ctx, t, errCh)
require.NoError(t, err)
_, ok := resp.Msg.(*TunnelMessage_Start)
require.True(t, ok)
@@ -194,7 +194,7 @@ func TestTunnel_PeerUpdate(t *testing.T) {
})
require.NoError(t, err)
// Then: the tunnel sends a PeerUpdate message
req := testutil.RequireRecvCtx(ctx, t, mgr.requests)
req := testutil.TryReceive(ctx, t, mgr.requests)
require.Nil(t, req.msg.Rpc)
require.NotNil(t, req.msg.GetPeerUpdate())
require.Len(t, req.msg.GetPeerUpdate().UpsertedWorkspaces, 1)
@@ -209,7 +209,7 @@ func TestTunnel_PeerUpdate(t *testing.T) {
errCh <- err
}()
// Then: a PeerUpdate message is sent using the Conn's state
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
require.NoError(t, err)
_, ok = resp.Msg.(*TunnelMessage_PeerUpdate)
require.True(t, ok)
@@ -243,8 +243,8 @@ func TestTunnel_NetworkSettings(t *testing.T) {
resp = r
errCh <- err
}()
testutil.RequireSendCtx(ctx, t, client.ch, conn)
err := testutil.RequireRecvCtx(ctx, t, errCh)
testutil.RequireSend(ctx, t, client.ch, conn)
err := testutil.TryReceive(ctx, t, errCh)
require.NoError(t, err)
_, ok := resp.Msg.(*TunnelMessage_Start)
require.True(t, ok)
@@ -257,11 +257,11 @@ func TestTunnel_NetworkSettings(t *testing.T) {
errCh <- err
}()
// Then: the tunnel sends a NetworkSettings message
req := testutil.RequireRecvCtx(ctx, t, mgr.requests)
req := testutil.TryReceive(ctx, t, mgr.requests)
require.NotNil(t, req.msg.Rpc)
require.Equal(t, uint32(1200), req.msg.GetNetworkSettings().Mtu)
go func() {
testutil.RequireSendCtx(ctx, t, mgr.sendCh, &ManagerMessage{
testutil.RequireSend(ctx, t, mgr.sendCh, &ManagerMessage{
Rpc: &RPC{ResponseTo: req.msg.Rpc.MsgId},
Msg: &ManagerMessage_NetworkSettings{
NetworkSettings: &NetworkSettingsResponse{
@@ -271,7 +271,7 @@ func TestTunnel_NetworkSettings(t *testing.T) {
})
}()
// And: `ApplyNetworkSettings` returns without error once the manager responds
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
require.NoError(t, err)
}
@@ -383,8 +383,8 @@ func TestTunnel_sendAgentUpdate(t *testing.T) {
resp = r
errCh <- err
}()
testutil.RequireSendCtx(ctx, t, client.ch, conn)
err := testutil.RequireRecvCtx(ctx, t, errCh)
testutil.RequireSend(ctx, t, client.ch, conn)
err := testutil.TryReceive(ctx, t, errCh)
require.NoError(t, err)
_, ok := resp.Msg.(*TunnelMessage_Start)
require.True(t, ok)
@@ -408,7 +408,7 @@ func TestTunnel_sendAgentUpdate(t *testing.T) {
},
})
require.NoError(t, err)
req := testutil.RequireRecvCtx(ctx, t, mgr.requests)
req := testutil.TryReceive(ctx, t, mgr.requests)
require.Nil(t, req.msg.Rpc)
require.NotNil(t, req.msg.GetPeerUpdate())
require.Len(t, req.msg.GetPeerUpdate().UpsertedAgents, 1)
@@ -420,7 +420,7 @@ func TestTunnel_sendAgentUpdate(t *testing.T) {
mClock.AdvanceNext()
// Then: the tunnel sends a PeerUpdate message of agent upserts,
// with the last handshake and latency set
req = testutil.RequireRecvCtx(ctx, t, mgr.requests)
req = testutil.TryReceive(ctx, t, mgr.requests)
require.Nil(t, req.msg.Rpc)
require.NotNil(t, req.msg.GetPeerUpdate())
require.Len(t, req.msg.GetPeerUpdate().UpsertedAgents, 1)
@@ -443,11 +443,11 @@ func TestTunnel_sendAgentUpdate(t *testing.T) {
},
})
require.NoError(t, err)
testutil.RequireRecvCtx(ctx, t, mgr.requests)
testutil.TryReceive(ctx, t, mgr.requests)
// The new update includes the new agent
mClock.AdvanceNext()
req = testutil.RequireRecvCtx(ctx, t, mgr.requests)
req = testutil.TryReceive(ctx, t, mgr.requests)
require.Nil(t, req.msg.Rpc)
require.NotNil(t, req.msg.GetPeerUpdate())
require.Len(t, req.msg.GetPeerUpdate().UpsertedAgents, 2)
@@ -474,11 +474,11 @@ func TestTunnel_sendAgentUpdate(t *testing.T) {
},
})
require.NoError(t, err)
testutil.RequireRecvCtx(ctx, t, mgr.requests)
testutil.TryReceive(ctx, t, mgr.requests)
// The new update doesn't include the deleted agent
mClock.AdvanceNext()
req = testutil.RequireRecvCtx(ctx, t, mgr.requests)
req = testutil.TryReceive(ctx, t, mgr.requests)
require.Nil(t, req.msg.Rpc)
require.NotNil(t, req.msg.GetPeerUpdate())
require.Len(t, req.msg.GetPeerUpdate().UpsertedAgents, 1)
@@ -506,9 +506,9 @@ func setupTunnel(t *testing.T, ctx context.Context, client *fakeClient, mClock q
mgr = manager
errCh <- err
}()
err := testutil.RequireRecvCtx(ctx, t, errCh)
err := testutil.TryReceive(ctx, t, errCh)
require.NoError(t, err)
err = testutil.RequireRecvCtx(ctx, t, errCh)
err = testutil.TryReceive(ctx, t, errCh)
require.NoError(t, err)
mgr.start()
return tun, mgr