gemini review + fix tests

Author: dishaprakash

tests/cloudsqlmysql/cloud_sql_mysql_mcp_test.go

@@ -106,32 +106,48 @@ func TestCloudSQLMySQLMCP(t *testing.T) {
 		if err != nil || statusCode != http.StatusOK || mcpResp.Result.IsError {
 			t.Fatalf("native error executing list_tables: %v", err)
 		}
-		got := mcpResp.Result.Content[0].Text
-		if !strings.Contains(got, tableNameParam) || !strings.Contains(got, tableNameAuth) {
-			t.Errorf("list_tables missing expected tables. Got: %s", got)
+		var gotTables string
+		for _, c := range mcpResp.Result.Content {
+			gotTables += c.Text
+		}
+		if !strings.Contains(gotTables, tableNameParam) || !strings.Contains(gotTables, tableNameAuth) {
+			t.Errorf("list_tables missing expected tables. Got: %s", gotTables)
 		}
 
 		go func() {
 			_ = pool.PingContext(ctx)
 			_, _ = pool.ExecContext(ctx, "SELECT sleep(5);")
 		}()
-		time.Sleep(1 * time.Second)
-
-		statusCode, mcpResp, err = tests.InvokeMCPTool(t, "list_active_queries", map[string]any{"min_duration_secs": 0}, nil)
-		if err != nil || statusCode != http.StatusOK || mcpResp.Result.IsError {
-			t.Fatalf("native error executing list_active_queries: %v", err)
+		var activeQueriesFound bool
+		for i := 0; i < 5; i++ {
+			time.Sleep(1 * time.Second)
+			statusCode, mcpResp, err = tests.InvokeMCPTool(t, "list_active_queries", map[string]any{"min_duration_secs": 0}, nil)
+			if err == nil && statusCode == http.StatusOK && !mcpResp.Result.IsError {
+				var gotQueries string
+				for _, c := range mcpResp.Result.Content {
+					gotQueries += c.Text
+				}
+				if strings.Contains(gotQueries, "SELECT sleep(5)") {
+					activeQueriesFound = true
+					break
+				}
+			}
 		}
-		if !strings.Contains(mcpResp.Result.Content[0].Text, "SELECT sleep(5)") {
-			t.Errorf("active queries did not contain test sleep query. Got: %s", mcpResp.Result.Content[0].Text)
+		if !activeQueriesFound {
+			t.Fatalf("active queries did not contain test sleep query after retries")
 		}
 
 		queryPlanSql := fmt.Sprintf("SELECT * FROM %s", tableNameParam)
 		statusCode, mcpResp, err = tests.InvokeMCPTool(t, "get_query_plan", map[string]any{"sql_statement": queryPlanSql}, nil)
 		if err != nil || statusCode != http.StatusOK || mcpResp.Result.IsError {
 			t.Fatalf("native error executing get_query_plan: %v", err)
 		}
-		if !strings.Contains(mcpResp.Result.Content[0].Text, "query_block") {
-			t.Errorf("query plan did not contain 'query_block'. Got: %s", mcpResp.Result.Content[0].Text)
+		var gotPlan string
+		for _, c := range mcpResp.Result.Content {
+			gotPlan += c.Text
+		}
+		if !strings.Contains(gotPlan, "query_block") {
+			t.Errorf("query plan did not contain 'query_block'. Got: %s", gotPlan)
 		}
 	})
 

tests/mcp_tool.go

@@ -177,7 +177,24 @@ func RunMCPCustomToolCallMethod(t *testing.T, toolName string, arguments map[str
 	if len(mcpResp.Result.Content) == 0 {
 		t.Fatalf("%s returned empty content field", toolName)
 	}
-	got := mcpResp.Result.Content[0].Text
+
+	// Gather all the text blocks
+	var blocks []string
+	for _, content := range mcpResp.Result.Content {
+		if content.Type == "text" {
+			blocks = append(blocks, strings.TrimSpace(content.Text))
+		}
+	}
+
+	var got string
+	// If the test is expecting a JSON array, format the stitched blocks as a JSON array
+	if want != "" && strings.HasPrefix(strings.TrimSpace(want), "[") {
+		got = "[" + strings.Join(blocks, ",") + "]"
+	} else {
+		// Otherwise, just concatenate them normally (for single responses or errors)
+		got = strings.Join(blocks, "")
+	}
+
 	if !strings.Contains(got, want) {
 		t.Fatalf(`expected %q to contain %q`, got, want)
 	}
@@ -278,7 +295,24 @@ func RunMCPToolInvokeTest(t *testing.T, select1Want string, options ...InvokeTes
 			if len(mcpResp.Result.Content) == 0 {
 				t.Fatalf("%s returned empty content field", tc.toolName)
 			}
-			got := mcpResp.Result.Content[0].Text
+
+			// Gather all the text blocks
+			var blocks []string
+			for _, content := range mcpResp.Result.Content {
+				if content.Type == "text" {
+					blocks = append(blocks, strings.TrimSpace(content.Text))
+				}
+			}
+
+			var got string
+			// If the test is expecting a JSON array, format the stitched blocks as a JSON array
+			if tc.wantResult != "" && strings.HasPrefix(strings.TrimSpace(tc.wantResult), "[") {
+				got = "[" + strings.Join(blocks, ",") + "]"
+			} else {
+				// Otherwise, just concatenate them normally (for single responses or errors)
+				got = strings.Join(blocks, "")
+			}
+
 			if !strings.Contains(got, tc.wantResult) {
 				t.Fatalf(`expected %q to contain %q`, got, tc.wantResult)
 			}
@@ -370,7 +404,7 @@ func GetExecuteSQLMCPExpectedTools() []MCPToolManifest {
 			Description: "Tool to execute sql",
 			InputSchema: map[string]any{
 				"type":       "object",
-				"properties": map[string]any{"sql": map[string]any{"type": "string", "description": "A valid SQL statement to execute."}},
+				"properties": map[string]any{"sql": map[string]any{"type": "string", "description": "The sql to execute."}},
 				"required":   []any{"sql"},
 			},
 		},
@@ -379,7 +413,7 @@ func GetExecuteSQLMCPExpectedTools() []MCPToolManifest {
 			Description: "Tool to execute sql",
 			InputSchema: map[string]any{
 				"type":       "object",
-				"properties": map[string]any{"sql": map[string]any{"type": "string", "description": "A valid SQL statement to execute."}},
+				"properties": map[string]any{"sql": map[string]any{"type": "string", "description": "The sql to execute."}},
 				"required":   []any{"sql"},
 			},
 		},