@@ -21,6 +21,7 @@ import (
2121 "net/http"
2222 "net/http/httptest"
2323 "net/url"
24+ "os"
2425 "reflect"
2526 "regexp"
2627 "sort"
@@ -33,6 +34,115 @@ import (
3334 "github.com/googleapis/genai-toolbox/tests"
3435)
3536
37+ var (
38+ AlloyDBProject = os .Getenv ("ALLOYDB_PROJECT" )
39+ AlloyDBLocation = os .Getenv ("ALLOYDB_REGION" )
40+ AlloyDBCluster = os .Getenv ("ALLOYDB_CLUSTER" )
41+ AlloyDBInstance = os .Getenv ("ALLOYDB_INSTANCE" )
42+ AlloyDBUser = os .Getenv ("ALLOYDB_POSTGRES_USER" )
43+ )
44+
45+ func getAlloyDBVars (t * testing.T ) map [string ]string {
46+ if AlloyDBProject == "" {
47+ t .Fatal ("'ALLOYDB_PROJECT' not set" )
48+ }
49+ if AlloyDBLocation == "" {
50+ t .Fatal ("'ALLOYDB_REGION' not set" )
51+ }
52+ if AlloyDBCluster == "" {
53+ t .Fatal ("'ALLOYDB_CLUSTER' not set" )
54+ }
55+ if AlloyDBInstance == "" {
56+ t .Fatal ("'ALLOYDB_INSTANCE' not set" )
57+ }
58+ if AlloyDBUser == "" {
59+ t .Fatal ("'ALLOYDB_USER' not set" )
60+ }
61+ return map [string ]string {
62+ "project" : AlloyDBProject ,
63+ "location" : AlloyDBLocation ,
64+ "cluster" : AlloyDBCluster ,
65+ "instance" : AlloyDBInstance ,
66+ "user" : AlloyDBUser ,
67+ }
68+ }
69+
70+ func getAlloyDBToolsConfig () map [string ]any {
71+ return map [string ]any {
72+ "sources" : map [string ]any {
73+ "alloydb-admin-source" : map [string ]any {
74+ "type" : "alloydb-admin" ,
75+ },
76+ },
77+ "tools" : map [string ]any {
78+ // Tool for RunAlloyDBToolGetTest
79+ "my-simple-tool" : map [string ]any {
80+ "type" : "alloydb-list-clusters" ,
81+ "source" : "alloydb-admin-source" ,
82+ "description" : "Simple tool to test end to end functionality." ,
83+ },
84+ // Tool for MCP test
85+ "my-param-tool" : map [string ]any {
86+ "type" : "alloydb-list-clusters" ,
87+ "source" : "alloydb-admin-source" ,
88+ "description" : "Tool to list clusters" ,
89+ },
90+ // Tool for MCP test that fails
91+ "my-fail-tool" : map [string ]any {
92+ "type" : "alloydb-list-clusters" ,
93+ "source" : "alloydb-admin-source" ,
94+ "description" : "Tool that will fail" ,
95+ },
96+ // AlloyDB specific tools
97+ "alloydb-list-clusters" : map [string ]any {
98+ "type" : "alloydb-list-clusters" ,
99+ "source" : "alloydb-admin-source" ,
100+ "description" : "Lists all AlloyDB clusters in a given project and location." ,
101+ },
102+ "alloydb-list-users" : map [string ]any {
103+ "type" : "alloydb-list-users" ,
104+ "source" : "alloydb-admin-source" ,
105+ "description" : "Lists all AlloyDB users within a specific cluster." ,
106+ },
107+ "alloydb-list-instances" : map [string ]any {
108+ "type" : "alloydb-list-instances" ,
109+ "source" : "alloydb-admin-source" ,
110+ "description" : "Lists all AlloyDB instances within a specific cluster." ,
111+ },
112+ "alloydb-get-cluster" : map [string ]any {
113+ "type" : "alloydb-get-cluster" ,
114+ "source" : "alloydb-admin-source" ,
115+ "description" : "Retrieves details of a specific AlloyDB cluster." ,
116+ },
117+ "alloydb-get-instance" : map [string ]any {
118+ "type" : "alloydb-get-instance" ,
119+ "source" : "alloydb-admin-source" ,
120+ "description" : "Retrieves details of a specific AlloyDB instance." ,
121+ },
122+ "alloydb-get-user" : map [string ]any {
123+ "type" : "alloydb-get-user" ,
124+ "source" : "alloydb-admin-source" ,
125+ "description" : "Retrieves details of a specific AlloyDB user." ,
126+ },
127+ "alloydb-create-cluster" : map [string ]any {
128+ "type" : "alloydb-create-cluster" ,
129+ "description" : "create cluster" ,
130+ "source" : "alloydb-admin-source" ,
131+ },
132+ "alloydb-create-instance" : map [string ]any {
133+ "type" : "alloydb-create-instance" ,
134+ "description" : "create instance" ,
135+ "source" : "alloydb-admin-source" ,
136+ },
137+ "alloydb-create-user" : map [string ]any {
138+ "type" : "alloydb-create-user" ,
139+ "description" : "create user" ,
140+ "source" : "alloydb-admin-source" ,
141+ },
142+ },
143+ }
144+ }
145+
36146func TestAlloyDBListTools (t * testing.T ) {
37147 ctx , cancel := context .WithTimeout (context .Background (), 2 * time .Minute )
38148 defer cancel ()
@@ -845,9 +955,8 @@ func TestAlloyDBCreateClusterMCP(t *testing.T) {
845955 ctx , cancel := context .WithTimeout (context .Background (), time .Minute )
846956 defer cancel ()
847957
848- args := []string {"--enable-api" }
849958 toolsFile := getAlloyDBToolsConfig ()
850- cmd , cleanupCmd , err := tests .StartCmd (ctx , toolsFile , args ... )
959+ cmd , cleanupCmd , err := tests .StartCmd (ctx , toolsFile )
851960 if err != nil {
852961 t .Fatalf ("command initialization returned an error: %v" , err )
853962 }
@@ -942,9 +1051,8 @@ func TestAlloyDBCreateInstanceMCP(t *testing.T) {
9421051 ctx , cancel := context .WithTimeout (context .Background (), time .Minute )
9431052 defer cancel ()
9441053
945- args := []string {"--enable-api" }
9461054 toolsFile := getAlloyDBToolsConfig ()
947- cmd , cleanupCmd , err := tests .StartCmd (ctx , toolsFile , args ... )
1055+ cmd , cleanupCmd , err := tests .StartCmd (ctx , toolsFile )
9481056 if err != nil {
9491057 t .Fatalf ("command initialization returned an error: %v" , err )
9501058 }
@@ -1049,9 +1157,8 @@ func TestAlloyDBCreateUserMCP(t *testing.T) {
10491157 ctx , cancel := context .WithTimeout (context .Background (), time .Minute )
10501158 defer cancel ()
10511159
1052- args := []string {"--enable-api" }
10531160 toolsFile := getAlloyDBToolsConfig ()
1054- cmd , cleanupCmd , err := tests .StartCmd (ctx , toolsFile , args ... )
1161+ cmd , cleanupCmd , err := tests .StartCmd (ctx , toolsFile )
10551162 if err != nil {
10561163 t .Fatalf ("command initialization returned an error: %v" , err )
10571164 }
0 commit comments