chore(coderd): fix test flake in TestAgentWebsocketMonitor_SendPings (#11518)

johnstcn · web-flow · commit 5ecb0db4f29c · 2024-01-10T08:45:46.000Z
diff --git a/coderd/workspaceagentsrpc_internal_test.go b/coderd/workspaceagentsrpc_internal_test.go
@@ -219,14 +219,21 @@ func TestAgentWebsocketMonitor_BuildOutdated(t *testing.T) {
 
 func TestAgentWebsocketMonitor_SendPings(t *testing.T) {
 	t.Parallel()
-	ctx := testutil.Context(t, testutil.WaitShort)
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+	t.Cleanup(cancel)
 	fConn := &fakePingerCloser{}
 	uut := &agentWebsocketMonitor{
 		pingPeriod: testutil.IntervalFast,
 		conn:       fConn,
 	}
-	go uut.sendPings(ctx)
+	done := make(chan struct{})
+	go func() {
+		uut.sendPings(ctx)
+		close(done)
+	}()
 	fConn.requireEventuallyHasPing(t)
+	cancel()
+	<-done
 	lastPing := uut.lastPing.Load()
 	require.NotNil(t, lastPing)
 }
