From 5ecb0db4f29c801623a1537cdb322d1831ed2daa Mon Sep 17 00:00:00 2001
From: Cian Johnston <cian@coder.com>
Date: Wed, 10 Jan 2024 08:45:46 +0000
Subject: [PATCH] chore(coderd): fix test flake in
 TestAgentWebsocketMonitor_SendPings (#11518)

---
 coderd/workspaceagentsrpc_internal_test.go | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/coderd/workspaceagentsrpc_internal_test.go b/coderd/workspaceagentsrpc_internal_test.go
index b8e6eb381c..834de4807d 100644
--- 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)
 }