Skip to content

Commit 79969dc

Browse files
anubhav756anubhav756
committed
test(source/alloydbainl): create MCP integration tests
1 parent 5346249 commit 79969dc

File tree

2 files changed

+282
-95
lines changed

2 files changed

+282
-95
lines changed

‎tests/alloydbainl/alloydb_ai_nl_integration_test.go‎

Lines changed: 0 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@ import (
2020
"encoding/json"
2121
"io"
2222
"net/http"
23-
"os"
2423
"reflect"
2524
"regexp"
2625
"strings"
@@ -32,47 +31,6 @@ import (
3231
"github.com/googleapis/genai-toolbox/tests"
3332
)
3433

35-
var (
36-
AlloyDBAINLSourceType = "alloydb-postgres"
37-
AlloyDBAINLToolType = "alloydb-ai-nl"
38-
AlloyDBAINLProject = os.Getenv("ALLOYDB_AI_NL_PROJECT")
39-
AlloyDBAINLRegion = os.Getenv("ALLOYDB_AI_NL_REGION")
40-
AlloyDBAINLCluster = os.Getenv("ALLOYDB_AI_NL_CLUSTER")
41-
AlloyDBAINLInstance = os.Getenv("ALLOYDB_AI_NL_INSTANCE")
42-
AlloyDBAINLDatabase = os.Getenv("ALLOYDB_AI_NL_DATABASE")
43-
AlloyDBAINLUser = os.Getenv("ALLOYDB_AI_NL_USER")
44-
AlloyDBAINLPass = os.Getenv("ALLOYDB_AI_NL_PASS")
45-
)
46-
47-
func getAlloyDBAINLVars(t *testing.T) map[string]any {
48-
switch "" {
49-
case AlloyDBAINLProject:
50-
t.Fatal("'ALLOYDB_AI_NL_PROJECT' not set")
51-
case AlloyDBAINLRegion:
52-
t.Fatal("'ALLOYDB_AI_NL_REGION' not set")
53-
case AlloyDBAINLCluster:
54-
t.Fatal("'ALLOYDB_AI_NL_CLUSTER' not set")
55-
case AlloyDBAINLInstance:
56-
t.Fatal("'ALLOYDB_AI_NL_INSTANCE' not set")
57-
case AlloyDBAINLDatabase:
58-
t.Fatal("'ALLOYDB_AI_NL_DATABASE' not set")
59-
case AlloyDBAINLUser:
60-
t.Fatal("'ALLOYDB_AI_NL_USER' not set")
61-
case AlloyDBAINLPass:
62-
t.Fatal("'ALLOYDB_AI_NL_PASS' not set")
63-
}
64-
return map[string]any{
65-
"type": AlloyDBAINLSourceType,
66-
"project": AlloyDBAINLProject,
67-
"cluster": AlloyDBAINLCluster,
68-
"instance": AlloyDBAINLInstance,
69-
"region": AlloyDBAINLRegion,
70-
"database": AlloyDBAINLDatabase,
71-
"user": AlloyDBAINLUser,
72-
"password": AlloyDBAINLPass,
73-
}
74-
}
75-
7634
func TestAlloyDBAINLToolEndpoints(t *testing.T) {
7735
sourceConfig := getAlloyDBAINLVars(t)
7836
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
@@ -277,59 +235,6 @@ func runAINLToolInvokeTest(t *testing.T) {
277235

278236
}
279237

280-
func getAINLToolsConfig(sourceConfig map[string]any) map[string]any {
281-
// Write config into a file and pass it to command
282-
toolsFile := map[string]any{
283-
"sources": map[string]any{
284-
"my-instance": sourceConfig,
285-
},
286-
"authServices": map[string]any{
287-
"my-google-auth": map[string]any{
288-
"type": "google",
289-
"clientId": tests.ClientId,
290-
},
291-
},
292-
"tools": map[string]any{
293-
"my-simple-tool": map[string]any{
294-
"type": AlloyDBAINLToolType,
295-
"source": "my-instance",
296-
"description": "Simple tool to test end to end functionality.",
297-
"nlConfig": "my_nl_config",
298-
},
299-
"my-auth-tool": map[string]any{
300-
"type": AlloyDBAINLToolType,
301-
"source": "my-instance",
302-
"description": "Tool to test authenticated parameters.",
303-
"nlConfig": "my_nl_config",
304-
"nlConfigParameters": []map[string]any{
305-
{
306-
"name": "email",
307-
"type": "string",
308-
"description": "user email",
309-
"authServices": []map[string]string{
310-
{
311-
"name": "my-google-auth",
312-
"field": "email",
313-
},
314-
},
315-
},
316-
},
317-
},
318-
"my-auth-required-tool": map[string]any{
319-
"type": AlloyDBAINLToolType,
320-
"source": "my-instance",
321-
"description": "Tool to test auth required invocation.",
322-
"nlConfig": "my_nl_config",
323-
"authRequired": []string{
324-
"my-google-auth",
325-
},
326-
},
327-
},
328-
}
329-
330-
return toolsFile
331-
}
332-
333238
func runAINLMCPToolCallMethod(t *testing.T) {
334239
sessionId := tests.RunInitialize(t, "2024-11-05")
335240
header := map[string]string{}

‎tests/alloydbainl/alloydb_ai_nl_mcp_test.go‎

Lines changed: 282 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,282 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package alloydbainl
16+
17+
import (
18+
"context"
19+
"net/http"
20+
"os"
21+
"regexp"
22+
"testing"
23+
"time"
24+
25+
"github.com/googleapis/genai-toolbox/internal/testutils"
26+
"github.com/googleapis/genai-toolbox/tests"
27+
)
28+
29+
func TestAlloyDBAINLToolEndpointsMCP(t *testing.T) {
30+
sourceConfig := getAlloyDBAINLVars(t)
31+
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
32+
defer cancel()
33+
34+
args := []string{"--enable-api"}
35+
36+
toolsFile := getAINLToolsConfig(sourceConfig)
37+
38+
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
39+
if err != nil {
40+
t.Fatalf("command initialization returned an error: %s", err)
41+
}
42+
defer cleanup()
43+
44+
waitCtx, cancelWait := context.WithTimeout(ctx, 10*time.Second)
45+
defer cancelWait()
46+
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
47+
if err != nil {
48+
t.Logf("toolbox command logs: \n%s", out)
49+
t.Fatalf("toolbox didn't start successfully: %s", err)
50+
}
51+
52+
runAINLToolGetMCPTest(t)
53+
runAINLToolInvokeMCPTest(t)
54+
}
55+
56+
func runAINLToolGetMCPTest(t *testing.T) {
57+
t.Run("list tools via MCP", func(t *testing.T) {
58+
statusCode, toolsList, err := tests.GetMCPToolsList(t, nil)
59+
if err != nil {
60+
t.Fatalf("native error executing tools/list: %s", err)
61+
}
62+
if statusCode != http.StatusOK {
63+
t.Fatalf("expected status 200, got %d", statusCode)
64+
}
65+
66+
// Verify that my-simple-tool is in the list
67+
found := false
68+
for _, tool := range toolsList {
69+
toolMap, ok := tool.(map[string]any)
70+
if !ok {
71+
continue
72+
}
73+
if toolMap["name"] == "my-simple-tool" {
74+
found = true
75+
break
76+
}
77+
}
78+
if !found {
79+
t.Errorf("expected tool 'my-simple-tool' not found in list")
80+
}
81+
})
82+
}
83+
84+
func runAINLToolInvokeMCPTest(t *testing.T) {
85+
idToken, err := tests.GetGoogleIdToken(tests.ClientId)
86+
if err != nil {
87+
t.Fatalf("error getting Google ID token: %s", err)
88+
}
89+
90+
invokeTcs := []struct {
91+
name string
92+
toolName string
93+
args map[string]any
94+
requestHeader map[string]string
95+
want string
96+
isErr bool
97+
wantStatusCode int
98+
}{
99+
{
100+
name: "invoke my-simple-tool",
101+
toolName: "my-simple-tool",
102+
args: map[string]any{"question": "return the number 1"},
103+
want: "{\"execute_nl_query\":{\"?column?\":1}}",
104+
isErr: false,
105+
},
106+
{
107+
name: "Invoke my-auth-tool with auth token",
108+
toolName: "my-auth-tool",
109+
args: map[string]any{"question": "can you show me the name of this user?"},
110+
requestHeader: map[string]string{"my-google-auth_token": idToken},
111+
want: "{\"execute_nl_query\":{\"name\":\"Alice\"}}",
112+
isErr: false,
113+
},
114+
{
115+
name: "Invoke my-auth-tool with invalid auth token",
116+
toolName: "my-auth-tool",
117+
args: map[string]any{"question": "return the number 1"},
118+
requestHeader: map[string]string{"my-google-auth_token": "INVALID_TOKEN"},
119+
isErr: true,
120+
},
121+
{
122+
name: "Invoke my-auth-tool without auth token",
123+
toolName: "my-auth-tool",
124+
args: map[string]any{"question": "return the number 1"},
125+
isErr: true,
126+
},
127+
{
128+
name: "Invoke my-auth-required-tool with auth token",
129+
toolName: "my-auth-required-tool",
130+
args: map[string]any{"question": "return the number 1"},
131+
requestHeader: map[string]string{"my-google-auth_token": idToken},
132+
isErr: false,
133+
want: "{\"execute_nl_query\":{\"?column?\":1}}",
134+
},
135+
{
136+
name: "Invoke my-auth-required-tool with invalid auth token",
137+
toolName: "my-auth-required-tool",
138+
args: map[string]any{"question": "return the number 1"},
139+
requestHeader: map[string]string{"my-google-auth_token": "INVALID_TOKEN"},
140+
isErr: true,
141+
wantStatusCode: 401,
142+
},
143+
{
144+
name: "Invoke my-auth-required-tool without auth token",
145+
toolName: "my-auth-required-tool",
146+
args: map[string]any{"question": "return the number 1"},
147+
isErr: true,
148+
wantStatusCode: 401,
149+
},
150+
}
151+
152+
for _, tc := range invokeTcs {
153+
t.Run(tc.name, func(t *testing.T) {
154+
statusCode, mcpResp, err := tests.InvokeMCPTool(t, tc.toolName, tc.args, tc.requestHeader)
155+
if err != nil {
156+
t.Fatalf("native error executing %s: %s", tc.toolName, err)
157+
}
158+
159+
expectedStatus := tc.wantStatusCode
160+
if expectedStatus == 0 {
161+
expectedStatus = http.StatusOK
162+
}
163+
if statusCode != expectedStatus {
164+
t.Fatalf("expected status %d, got %d", expectedStatus, statusCode)
165+
}
166+
167+
if tc.isErr {
168+
if mcpResp.Error == nil && !mcpResp.Result.IsError {
169+
t.Fatalf("expected error result or JSON-RPC error, got success")
170+
}
171+
} else {
172+
if mcpResp.Error != nil {
173+
t.Fatalf("expected success, got JSON-RPC error: %v", mcpResp.Error)
174+
}
175+
if mcpResp.Result.IsError {
176+
t.Fatalf("expected success result, got tool error: %v", mcpResp.Result)
177+
}
178+
if len(mcpResp.Result.Content) == 0 {
179+
t.Fatalf("expected at least one content item, got none")
180+
}
181+
got := mcpResp.Result.Content[0].Text
182+
if got != tc.want {
183+
t.Fatalf("unexpected value: got %q, want %q", got, tc.want)
184+
}
185+
}
186+
})
187+
}
188+
}
189+
190+
var (
191+
AlloyDBAINLSourceType = "alloydb-postgres"
192+
AlloyDBAINLToolType = "alloydb-ai-nl"
193+
AlloyDBAINLProject = os.Getenv("ALLOYDB_AI_NL_PROJECT")
194+
AlloyDBAINLRegion = os.Getenv("ALLOYDB_AI_NL_REGION")
195+
AlloyDBAINLCluster = os.Getenv("ALLOYDB_AI_NL_CLUSTER")
196+
AlloyDBAINLInstance = os.Getenv("ALLOYDB_AI_NL_INSTANCE")
197+
AlloyDBAINLDatabase = os.Getenv("ALLOYDB_AI_NL_DATABASE")
198+
AlloyDBAINLUser = os.Getenv("ALLOYDB_AI_NL_USER")
199+
AlloyDBAINLPass = os.Getenv("ALLOYDB_AI_NL_PASS")
200+
)
201+
202+
func getAlloyDBAINLVars(t *testing.T) map[string]any {
203+
switch "" {
204+
case AlloyDBAINLProject:
205+
t.Fatal("'ALLOYDB_AI_NL_PROJECT' not set")
206+
case AlloyDBAINLRegion:
207+
t.Fatal("'ALLOYDB_AI_NL_REGION' not set")
208+
case AlloyDBAINLCluster:
209+
t.Fatal("'ALLOYDB_AI_NL_CLUSTER' not set")
210+
case AlloyDBAINLInstance:
211+
t.Fatal("'ALLOYDB_AI_NL_INSTANCE' not set")
212+
case AlloyDBAINLDatabase:
213+
t.Fatal("'ALLOYDB_AI_NL_DATABASE' not set")
214+
case AlloyDBAINLUser:
215+
t.Fatal("'ALLOYDB_AI_NL_USER' not set")
216+
case AlloyDBAINLPass:
217+
t.Fatal("'ALLOYDB_AI_NL_PASS' not set")
218+
}
219+
return map[string]any{
220+
"type": AlloyDBAINLSourceType,
221+
"project": AlloyDBAINLProject,
222+
"cluster": AlloyDBAINLCluster,
223+
"instance": AlloyDBAINLInstance,
224+
"region": AlloyDBAINLRegion,
225+
"database": AlloyDBAINLDatabase,
226+
"user": AlloyDBAINLUser,
227+
"password": AlloyDBAINLPass,
228+
}
229+
}
230+
231+
func getAINLToolsConfig(sourceConfig map[string]any) map[string]any {
232+
// Write config into a file and pass it to command
233+
toolsFile := map[string]any{
234+
"sources": map[string]any{
235+
"my-instance": sourceConfig,
236+
},
237+
"authServices": map[string]any{
238+
"my-google-auth": map[string]any{
239+
"type": "google",
240+
"clientId": tests.ClientId,
241+
},
242+
},
243+
"tools": map[string]any{
244+
"my-simple-tool": map[string]any{
245+
"type": AlloyDBAINLToolType,
246+
"source": "my-instance",
247+
"description": "Simple tool to test end to end functionality.",
248+
"nlConfig": "my_nl_config",
249+
},
250+
"my-auth-tool": map[string]any{
251+
"type": AlloyDBAINLToolType,
252+
"source": "my-instance",
253+
"description": "Tool to test authenticated parameters.",
254+
"nlConfig": "my_nl_config",
255+
"nlConfigParameters": []map[string]any{
256+
{
257+
"name": "email",
258+
"type": "string",
259+
"description": "user email",
260+
"authServices": []map[string]string{
261+
{
262+
"name": "my-google-auth",
263+
"field": "email",
264+
},
265+
},
266+
},
267+
},
268+
},
269+
"my-auth-required-tool": map[string]any{
270+
"type": AlloyDBAINLToolType,
271+
"source": "my-instance",
272+
"description": "Tool to test auth required invocation.",
273+
"nlConfig": "my_nl_config",
274+
"authRequired": []string{
275+
"my-google-auth",
276+
},
277+
},
278+
},
279+
}
280+
281+
return toolsFile
282+
}

0 commit comments

Comments
 (0)