fix: fix flake in TestWorkspaceAgent_Metadata_CatchMemoryLeak (#13553)

spikecurtis · web-flow · commit ba7d1835e5c5 · 2024-06-12T18:33:22.000+04:00
Fixes flake seen here: https://github.com/coder/coder/actions/runs/9461246505/job/26061605278 #13486 subtly changes the test so that `post` uses the new v2 Agent API, and when canceling context, there is a race condition where the yamux session underpinning the API can get torn down before the RPC processes the canceled context, yielding a different error response than the test was previously expecting. I've refactored the test to just stop posting when the test finishes, rather than depend on a context cancel to end the posting goroutine.
diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go
@@ -1350,7 +1350,7 @@ func TestWorkspaceAgent_Metadata_CatchMemoryLeak(t *testing.T) {
 	agentClient := agentsdk.New(client.URL)
 	agentClient.SetSessionToken(r.AgentToken)
 
-	ctx, cancel := context.WithCancel(testutil.Context(t, testutil.WaitSuperLong))
+	ctx := testutil.Context(t, testutil.WaitSuperLong)
 	conn, err := agentClient.ConnectRPC(ctx)
 	require.NoError(t, err)
 	defer func() {
@@ -1404,20 +1404,21 @@ func TestWorkspaceAgent_Metadata_CatchMemoryLeak(t *testing.T) {
 
 	postDone := testutil.Go(t, func() {
 		for {
+			select {
+			case <-metadataDone:
+				return
+			default:
+			}
 			// We need to send two separate metadata updates to trigger the
 			// memory leak. foo2 will cause the number of foo1 to be doubled, etc.
-			err = post(ctx, "foo1", "hi")
+			err := post(ctx, "foo1", "hi")
 			if err != nil {
-				if !xerrors.Is(err, context.Canceled) {
-					assert.NoError(t, err, "post metadata foo1")
-				}
+				assert.NoError(t, err, "post metadata foo1")
 				return
 			}
 			err = post(ctx, "foo2", "bye")
 			if err != nil {
-				if !xerrors.Is(err, context.Canceled) {
-					assert.NoError(t, err, "post metadata foo1")
-				}
+				assert.NoError(t, err, "post metadata foo1")
 				return
 			}
 		}
@@ -1436,13 +1437,8 @@ func TestWorkspaceAgent_Metadata_CatchMemoryLeak(t *testing.T) {
 	// testing it is not straightforward.
 	db.err.Store(&wantErr)
 
-	select {
-	case <-ctx.Done():
-		t.Fatal("timeout waiting for SSE to close")
-	case <-metadataDone:
-	}
-	cancel()
-	<-postDone
+	testutil.RequireRecvCtx(ctx, t, metadataDone)
+	testutil.RequireRecvCtx(ctx, t, postDone)
 }
 
 func TestWorkspaceAgent_Startup(t *testing.T) {
