Skip to content

Commit 0a48e1c

Browse files
dishaprakashdishaprakash
committed
fix
1 parent 1a0596d commit 0a48e1c

File tree

4 files changed

+169
-205
lines changed

4 files changed

+169
-205
lines changed

‎tests/cloudsqlmysql/cloud_sql_mysql_mcp_test.go‎

Lines changed: 7 additions & 154 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
1-
// Copyright 2026 Google LLC
1+
// Copyright 2024 Google LLC
22
//
33
// Licensed under the Apache License, Version 2.0 (the "License");
44
// you may not use this file except in compliance with the License.
55
// You may obtain a copy of the License at
66
//
7-
// http://www.apache.org/licenses/LICENSE-2.0
7+
// http://www.apache.org/licenses/LICENSE-2.0
88
//
99
// Unless required by applicable law or agreed to in writing, software
1010
// distributed under the License is distributed on an "AS IS" BASIS,
@@ -16,8 +16,6 @@ package cloudsqlmysql
1616

1717
import (
1818
"context"
19-
"fmt"
20-
"net/http"
2119
"regexp"
2220
"strings"
2321
"testing"
@@ -28,7 +26,7 @@ import (
2826
"github.com/googleapis/genai-toolbox/tests"
2927
)
3028

31-
func TestCloudSQLMySQLToolEndpointsMCP(t *testing.T) {
29+
func TestCloudSQLMySQLMCPToolEndpoints(t *testing.T) {
3230
sourceConfig := getCloudSQLMySQLVars(t)
3331
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
3432
defer cancel()
@@ -176,159 +174,14 @@ func TestCloudSQLMySQLToolEndpointsMCP(t *testing.T) {
176174
tests.RunMCPToolsListMethod(t, expectedTools)
177175
})
178176

179-
// Run standard shared tests (MCP variants)
180-
tests.RunMCPToolInvokeTest(t, select1Want, tests.DisableArrayTest())
177+
tests.RunToolInvokeTest(t, select1Want, tests.DisableArrayTest(), tests.WithMCPInvoke(true))
181178
tests.RunMCPToolCallMethod(t, mcpMyFailToolWant, mcpSelect1Want)
179+
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want, tests.WithMCPExecuteSql(true))
180+
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.WithMCPTemplate()) // Assuming generic WithMCP for this helper
182181

183-
// Execute SQL test
184-
statusCode, mcpResp, err := tests.InvokeMCPTool(t, "my-exec-sql-tool", map[string]any{"sql": createTableStatement}, nil)
185-
if err == nil && statusCode == http.StatusOK && !mcpResp.Result.IsError {
186-
tests.RunMCPCustomToolCallMethod(t, "my-exec-sql-tool", map[string]any{"sql": "SELECT 1"}, select1Want)
187-
}
188-
189-
// Run template parameter test
190-
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.WithMCPTemplate())
191-
192-
// Run specific MySQL tool tests
182+
// Run specific MySQL tool tests over MCP
193183
const expectedOwner = "'toolbox-identity'@'%'"
194184
tests.RunMySQLListTablesTest(t, CloudSQLMySQLDatabase, tableNameParam, tableNameAuth, expectedOwner, tests.WithMCPExec())
195185
tests.RunMySQLListActiveQueriesTest(t, ctx, pool, tests.WithMCPExec())
196186
tests.RunMySQLGetQueryPlanTest(t, ctx, pool, CloudSQLMySQLDatabase, tableNameParam, tests.WithMCPExec())
197-
198-
// Ensure MCP validation
199-
t.Run("verify parameter validation for prebuilt tools", func(t *testing.T) {
200-
statusCode, mcpResp, err := tests.InvokeMCPTool(t, "get_query_plan", map[string]any{}, nil)
201-
if err != nil && statusCode == http.StatusOK {
202-
t.Fatalf("native error executing get_query_plan: %v", err)
203-
}
204-
if statusCode != http.StatusOK {
205-
t.Fatalf("expected status 200, got %d", statusCode)
206-
}
207-
tests.AssertMCPError(t, mcpResp, `parameter "sql_statement" is required`)
208-
})
209-
}
210-
211-
// Test connection with different IP type over MCP
212-
func TestCloudSQLMySQLMCPIpConnection(t *testing.T) {
213-
sourceConfig := getCloudSQLMySQLVars(t)
214-
tcs := []struct {
215-
name string
216-
ipType string
217-
}{
218-
{
219-
name: "public ip",
220-
ipType: "public",
221-
},
222-
{
223-
name: "private ip",
224-
ipType: "private",
225-
},
226-
}
227-
for _, tc := range tcs {
228-
t.Run(tc.name, func(t *testing.T) {
229-
sourceConfig["ipType"] = tc.ipType
230-
err := tests.RunSourceConnectionTest(t, sourceConfig, CloudSQLMySQLToolType, tests.WithMCP())
231-
if err != nil {
232-
t.Fatalf("Connection test failure via MCP: %s", err)
233-
}
234-
})
235-
}
236-
}
237-
238-
func TestCloudSQLMySQLMCPIAMConnection(t *testing.T) {
239-
getCloudSQLMySQLVars(t)
240-
serviceAccountEmail, _, _ := strings.Cut(tests.ServiceAccountEmail, "@")
241-
242-
noPassSourceConfig := map[string]any{
243-
"type": CloudSQLMySQLSourceType,
244-
"project": CloudSQLMySQLProject,
245-
"instance": CloudSQLMySQLInstance,
246-
"region": CloudSQLMySQLRegion,
247-
"database": CloudSQLMySQLDatabase,
248-
"user": serviceAccountEmail,
249-
}
250-
noUserSourceConfig := map[string]any{
251-
"type": CloudSQLMySQLSourceType,
252-
"project": CloudSQLMySQLProject,
253-
"instance": CloudSQLMySQLInstance,
254-
"region": CloudSQLMySQLRegion,
255-
"database": CloudSQLMySQLDatabase,
256-
"password": "random",
257-
}
258-
noUserNoPassSourceConfig := map[string]any{
259-
"type": CloudSQLMySQLSourceType,
260-
"project": CloudSQLMySQLProject,
261-
"instance": CloudSQLMySQLInstance,
262-
"region": CloudSQLMySQLRegion,
263-
"database": CloudSQLMySQLDatabase,
264-
}
265-
tcs := []struct {
266-
name string
267-
sourceConfig map[string]any
268-
isErr bool
269-
}{
270-
{
271-
name: "no user no pass",
272-
sourceConfig: noUserNoPassSourceConfig,
273-
isErr: false,
274-
},
275-
{
276-
name: "no password",
277-
sourceConfig: noPassSourceConfig,
278-
isErr: false,
279-
},
280-
{
281-
name: "no user",
282-
sourceConfig: noUserSourceConfig,
283-
isErr: true,
284-
},
285-
}
286-
for i, tc := range tcs {
287-
t.Run(tc.name, func(t *testing.T) {
288-
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
289-
defer cancel()
290-
291-
uniqueSourceName := fmt.Sprintf("iam-test-mcp-%d", i)
292-
293-
toolsFile := map[string]any{
294-
"sources": map[string]any{
295-
uniqueSourceName: tc.sourceConfig,
296-
},
297-
"tools": map[string]any{
298-
"my-simple-tool": map[string]any{
299-
"type": CloudSQLMySQLToolType,
300-
"source": uniqueSourceName,
301-
"description": "Simple tool to test end to end functionality.",
302-
"statement": "SELECT 1;",
303-
},
304-
},
305-
}
306-
307-
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile)
308-
if err != nil {
309-
t.Fatalf("command initialization returned an error: %s", err)
310-
}
311-
defer cleanup()
312-
313-
waitCtx, waitCancel := context.WithTimeout(ctx, 10*time.Second)
314-
defer waitCancel()
315-
316-
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
317-
if err != nil {
318-
if tc.isErr {
319-
outLower := strings.ToLower(out)
320-
if !strings.Contains(outLower, "error") && !strings.Contains(outLower, "fail") {
321-
t.Fatalf("Expected an authentication/connection error, but server failed for another reason. Output: %s", out)
322-
}
323-
return
324-
}
325-
t.Logf("toolbox command logs: \n%s", out)
326-
t.Fatalf("Connection test failure: toolbox didn't start successfully: %s", err)
327-
}
328-
329-
if tc.isErr {
330-
t.Fatalf("Expected error but test passed.")
331-
}
332-
})
333-
}
334187
}

‎tests/option.go‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,23 @@
1414

1515
package tests
1616

17+
/* Configurations for RunSourceConnectionTest() */
18+
19+
// ConnectionTestConfig holds configuration for source connection tests
20+
type ConnectionTestConfig struct {
21+
UseMCP bool
22+
}
23+
24+
// ConnectionTestOption configures how the source connection test runs
25+
type ConnectionTestOption func(*ConnectionTestConfig)
26+
27+
// WithMCP configures the connection test to verify via the MCP protocol instead of the legacy API
28+
func WithMCP() ConnectionTestOption {
29+
return func(c *ConnectionTestConfig) {
30+
c.UseMCP = true
31+
}
32+
}
33+
1734
/* Configurations for RunToolInvokeTest() */
1835

1936
// InvokeTestConfig represents the various configuration options for RunToolInvokeTest()
@@ -28,10 +45,18 @@ type InvokeTestConfig struct {
2845
supportArrayParam bool
2946
supportClientAuth bool
3047
supportSelect1Auth bool
48+
IsMCP bool
3149
}
3250

3351
type InvokeTestOption func(*InvokeTestConfig)
3452

53+
// WithMCPInvoke enables the MCP routing for standard Tool Invoke tests
54+
func WithMCPInvoke(isMCP bool) InvokeTestOption {
55+
return func(c *InvokeTestConfig) {
56+
c.IsMCP = isMCP
57+
}
58+
}
59+
3560
// WithMyAuthToolWant represents the response value for my-auth-tool.
3661
// e.g. tests.RunToolInvokeTest(t, select1Want, tests.WithMyAuthToolWant("custom"))
3762
func WithMyAuthToolWant(s string) InvokeTestOption {
@@ -164,6 +189,7 @@ type ExecuteSqlTestConfig struct {
164189
createWant string
165190
dropWant string
166191
selectEmptyWant string
192+
IsMCP bool
167193
}
168194

169195
type ExecuteSqlOption func(*ExecuteSqlTestConfig)
@@ -176,6 +202,13 @@ func WithSelect1Statement(s string) ExecuteSqlOption {
176202
}
177203
}
178204

205+
// WithMCPExecuteSql enables the MCP routing for ExecuteSql tests
206+
func WithMCPExecuteSql(isMCP bool) ExecuteSqlOption {
207+
return func(c *ExecuteSqlTestConfig) {
208+
c.IsMCP = isMCP
209+
}
210+
}
211+
179212
// WithExecuteCreateWant represents the expected response for a CREATE TABLE statement.
180213
func WithExecuteCreateWant(s string) ExecuteSqlOption {
181214
return func(c *ExecuteSqlTestConfig) {

‎tests/source.go‎

Lines changed: 0 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -29,21 +29,6 @@ import (
2929
"github.com/googleapis/genai-toolbox/internal/testutils"
3030
)
3131

32-
// ConnectionTestConfig holds configuration for source connection tests
33-
type ConnectionTestConfig struct {
34-
UseMCP bool
35-
}
36-
37-
// ConnectionTestOption configures how the source connection test runs
38-
type ConnectionTestOption func(*ConnectionTestConfig)
39-
40-
// WithMCP configures the connection test to verify via the MCP protocol instead of the legacy API
41-
func WithMCP() ConnectionTestOption {
42-
return func(c *ConnectionTestConfig) {
43-
c.UseMCP = true
44-
}
45-
}
46-
4732
// RunSourceConnection test for source connection
4833
func RunSourceConnectionTest(t *testing.T, sourceConfig map[string]any, toolType string, options ...ConnectionTestOption) error {
4934
configs := &ConnectionTestConfig{

0 commit comments

Comments
 (0)