fix tests

johnstcn · johnstcn · commit 86d66b1a9cf7 · 2025-04-25T15:03:15.000+01:00
diff --git a/cli/exp_mcp_test.go b/cli/exp_mcp_test.go
@@ -33,10 +33,6 @@ func TestExpMcpServer(t *testing.T) {
 		ctx := testutil.Context(t, testutil.WaitShort)
 		cmdDone := make(chan struct{})
 		cancelCtx, cancel := context.WithCancel(ctx)
-		t.Cleanup(func() {
-			cancel()
-			<-cmdDone
-		})
 
 		// Given: a running coder deployment
 		client := coderdtest.New(t, nil)
@@ -93,6 +89,8 @@ func TestExpMcpServer(t *testing.T) {
 		require.NoError(t, err, "should have received a valid JSON response from the tool")
 		// Ensure the tool returns the expected user
 		require.Contains(t, output, owner.UserID.String(), "should have received the expected user ID")
+		cancel()
+		<-cmdDone
 	})
 
 	t.Run("OK", func(t *testing.T) {
diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go
@@ -341,7 +341,6 @@ func TestWorkspaceAgentLogs(t *testing.T) {
 
 func TestWorkspaceAgentAppStatus(t *testing.T) {
 	t.Parallel()
-	ctx := testutil.Context(t, testutil.WaitMedium)
 	client, db := coderdtest.NewWithDatabase(t, nil)
 	user := coderdtest.CreateFirstUser(t, client)
 	client, user2 := coderdtest.CreateAnotherUser(t, client, user.OrganizationID)
@@ -362,6 +361,7 @@ func TestWorkspaceAgentAppStatus(t *testing.T) {
 	agentClient.SetSessionToken(r.AgentToken)
 	t.Run("Success", func(t *testing.T) {
 		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitShort)
 		err := agentClient.PatchAppStatus(ctx, agentsdk.PatchAppStatus{
 			AppSlug: "vscode",
 			Message: "testing",
@@ -385,6 +385,7 @@ func TestWorkspaceAgentAppStatus(t *testing.T) {
 
 	t.Run("FailUnknownApp", func(t *testing.T) {
 		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitShort)
 		err := agentClient.PatchAppStatus(ctx, agentsdk.PatchAppStatus{
 			AppSlug: "unknown",
 			Message: "testing",
@@ -399,6 +400,7 @@ func TestWorkspaceAgentAppStatus(t *testing.T) {
 
 	t.Run("FailUnknownState", func(t *testing.T) {
 		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitShort)
 		err := agentClient.PatchAppStatus(ctx, agentsdk.PatchAppStatus{
 			AppSlug: "vscode",
 			Message: "testing",
@@ -413,6 +415,7 @@ func TestWorkspaceAgentAppStatus(t *testing.T) {
 
 	t.Run("FailTooLong", func(t *testing.T) {
 		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitShort)
 		err := agentClient.PatchAppStatus(ctx, agentsdk.PatchAppStatus{
 			AppSlug: "vscode",
 			Message: strings.Repeat("a", 161),
