fix tests

johnstcn · johnstcn · commit 0f1d4acdf674 · 2025-04-29T14:25:12.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),
diff --git a/codersdk/toolsdk/toolsdk.go b/codersdk/toolsdk/toolsdk.go
@@ -156,6 +156,13 @@ func WithCleanContext(h GenericHandlerFunc) GenericHandlerFunc {
 	return func(parent context.Context, tb Deps, args json.RawMessage) (ret json.RawMessage, err error) {
 		child, childCancel := context.WithCancel(context.Background())
 		defer childCancel()
+		// Ensure that the child context has the same deadline as the parent
+		// context.
+		if deadline, ok := parent.Deadline(); ok {
+			deadlineCtx, deadlineCancel := context.WithDeadline(child, deadline)
+			defer deadlineCancel()
+			child = deadlineCtx
+		}
 		// Ensure that cancellation propagates from the parent context to the child context.
 		go func() {
 			select {
@@ -165,13 +172,6 @@ func WithCleanContext(h GenericHandlerFunc) GenericHandlerFunc {
 				childCancel()
 			}
 		}()
-		// Also ensure that the child context has the same deadline as the parent
-		// context.
-		if deadline, ok := parent.Deadline(); ok {
-			deadlineCtx, deadlineCancel := context.WithDeadline(child, deadline)
-			defer deadlineCancel()
-			child = deadlineCtx
-		}
 		return h(child, tb, args)
 	}
 }
