chore: import upstream snapshot with attribution

tests/cloudsqlpg/cloud_sql_pg_create_instances_test.go

@@ -0,0 +1,278 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package cloudsqlpg
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"reflect"
+	"regexp"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/googleapis/mcp-toolbox/internal/testutils"
+	"github.com/googleapis/mcp-toolbox/tests"
+	"google.golang.org/api/sqladmin/v1"
+)
+
+var (
+	createInstanceToolType = "cloud-sql-postgres-create-instance"
+)
+
+type createInstanceTransport struct {
+	transport http.RoundTripper
+	url       *url.URL
+}
+
+func (t *createInstanceTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	if strings.HasPrefix(req.URL.String(), "https://sqladmin.googleapis.com") {
+		req.URL.Scheme = t.url.Scheme
+		req.URL.Host = t.url.Host
+	}
+	return t.transport.RoundTrip(req)
+}
+
+type masterHandler struct {
+	t *testing.T
+}
+
+func (h *masterHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	ua := r.Header.Get("User-Agent")
+	if !strings.Contains(ua, "genai-toolbox/") {
+		h.t.Errorf("User-Agent header not found in %q", ua)
+	}
+
+	var body sqladmin.DatabaseInstance
+	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+		h.t.Fatalf("failed to decode request body: %v", err)
+	}
+
+	instanceName := body.Name
+	if instanceName == "" {
+		http.Error(w, "missing instance name", http.StatusBadRequest)
+		return
+	}
+
+	var expectedBody sqladmin.DatabaseInstance
+	var response any
+	var statusCode int
+
+	switch instanceName {
+	case "instance1":
+		expectedBody = sqladmin.DatabaseInstance{
+			Project:         "p1",
+			Name:            "instance1",
+			DatabaseVersion: "POSTGRES_15",
+			RootPassword:    "password123",
+			Settings: &sqladmin.Settings{
+				AvailabilityType: "REGIONAL",
+				Edition:          "ENTERPRISE_PLUS",
+				Tier:             "db-perf-optimized-N-8",
+				DataDiskSizeGb:   250,
+				DataDiskType:     "PD_SSD",
+			},
+		}
+		response = map[string]any{"name": "op1", "status": "PENDING"}
+		statusCode = http.StatusOK
+	case "instance2":
+		expectedBody = sqladmin.DatabaseInstance{
+			Project:         "p2",
+			Name:            "instance2",
+			DatabaseVersion: "POSTGRES_17",
+			RootPassword:    "password456",
+			Settings: &sqladmin.Settings{
+				AvailabilityType: "ZONAL",
+				Edition:          "ENTERPRISE_PLUS",
+				Tier:             "db-perf-optimized-N-2",
+				DataDiskSizeGb:   100,
+				DataDiskType:     "PD_SSD",
+			},
+		}
+		response = map[string]any{"name": "op2", "status": "RUNNING"}
+		statusCode = http.StatusOK
+	default:
+		http.Error(w, fmt.Sprintf("unhandled instance name: %s", instanceName), http.StatusInternalServerError)
+		return
+	}
+
+	if expectedBody.Project != body.Project {
+		h.t.Errorf("unexpected project: got %q, want %q", body.Project, expectedBody.Project)
+	}
+	if expectedBody.Name != body.Name {
+		h.t.Errorf("unexpected name: got %q, want %q", body.Name, expectedBody.Name)
+	}
+	if expectedBody.DatabaseVersion != body.DatabaseVersion {
+		h.t.Errorf("unexpected databaseVersion: got %q, want %q", body.DatabaseVersion, expectedBody.DatabaseVersion)
+	}
+	if expectedBody.RootPassword != body.RootPassword {
+		h.t.Errorf("unexpected rootPassword: got %q, want %q", body.RootPassword, expectedBody.RootPassword)
+	}
+	if diff := cmp.Diff(expectedBody.Settings, body.Settings); diff != "" {
+		h.t.Errorf("unexpected request body settings (-want +got):\n%s", diff)
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(statusCode)
+	if err := json.NewEncoder(w).Encode(response); err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+	}
+}
+
+func TestCreateInstanceToolEndpoints(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+
+	handler := &masterHandler{t: t}
+	server := httptest.NewServer(handler)
+	defer server.Close()
+
+	serverURL, err := url.Parse(server.URL)
+	if err != nil {
+		t.Fatalf("failed to parse server URL: %v", err)
+	}
+
+	originalTransport := http.DefaultClient.Transport
+	if originalTransport == nil {
+		originalTransport = http.DefaultTransport
+	}
+	http.DefaultClient.Transport = &createInstanceTransport{
+		transport: originalTransport,
+		url:       serverURL,
+	}
+	t.Cleanup(func() {
+		http.DefaultClient.Transport = originalTransport
+	})
+
+	args := []string{"--enable-api"}
+	toolsFile := getCreateInstanceToolsConfig()
+	cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
+	if err != nil {
+		t.Fatalf("command initialization returned an error: %s", err)
+	}
+	defer cleanup()
+
+	waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+	out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
+	if err != nil {
+		t.Logf("toolbox command logs: \n%s", out)
+		t.Fatalf("toolbox didn't start successfully: %s", err)
+	}
+
+	tcs := []struct {
+		name        string
+		toolName    string
+		body        string
+		want        string
+		expectError bool
+		errorStatus int
+	}{
+		{
+			name:     "successful creation - production",
+			toolName: "create-instance-prod",
+			body:     `{"project": "p1", "name": "instance1", "databaseVersion": "POSTGRES_15", "rootPassword": "password123", "editionPreset": "Production"}`,
+			want:     `{"name":"op1","status":"PENDING"}`,
+		},
+		{
+			name:     "successful creation - development",
+			toolName: "create-instance-dev",
+			body:     `{"project": "p2", "name": "instance2", "rootPassword": "password456", "editionPreset": "Development"}`,
+			want:     `{"name":"op2","status":"RUNNING"}`,
+		},
+		{
+			name:     "missing required parameter",
+			toolName: "create-instance-prod",
+			body:     `{"name": "instance1"}`,
+			want:     `{"error":"parameter \"project\" is required"}`,
+		},
+	}
+
+	for _, tc := range tcs {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			api := fmt.Sprintf("http://127.0.0.1:5000/api/tool/%s/invoke", tc.toolName)
+			req, err := http.NewRequest(http.MethodPost, api, bytes.NewBufferString(tc.body))
+			if err != nil {
+				t.Fatalf("unable to create request: %s", err)
+			}
+			req.Header.Add("Content-type", "application/json")
+			resp, err := http.DefaultClient.Do(req)
+			if err != nil {
+				t.Fatalf("unable to send request: %s", err)
+			}
+			defer resp.Body.Close()
+
+			if tc.expectError {
+				if resp.StatusCode != tc.errorStatus {
+					bodyBytes, _ := io.ReadAll(resp.Body)
+					t.Fatalf("expected status %d but got %d: %s", tc.errorStatus, resp.StatusCode, string(bodyBytes))
+				}
+				return
+			}
+
+			if resp.StatusCode != http.StatusOK {
+				bodyBytes, _ := io.ReadAll(resp.Body)
+				t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes))
+			}
+
+			var result struct {
+				Result string `json:"result"`
+			}
+			if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
+				t.Fatalf("failed to decode response: %v", err)
+			}
+
+			var got, want map[string]any
+			if err := json.Unmarshal([]byte(result.Result), &got); err != nil {
+				t.Fatalf("failed to unmarshal result: %v", err)
+			}
+			if err := json.Unmarshal([]byte(tc.want), &want); err != nil {
+				t.Fatalf("failed to unmarshal want: %v", err)
+			}
+
+			if !reflect.DeepEqual(got, want) {
+				t.Fatalf("unexpected result: got %+v, want %+v", got, want)
+			}
+		})
+	}
+}
+
+func getCreateInstanceToolsConfig() map[string]any {
+	return map[string]any{
+		"sources": map[string]any{
+			"my-cloud-sql-source": map[string]any{
+				"type": "cloud-sql-admin",
+			},
+		},
+		"tools": map[string]any{
+			"create-instance-prod": map[string]any{
+				"type":   createInstanceToolType,
+				"source": "my-cloud-sql-source",
+			},
+			"create-instance-dev": map[string]any{
+				"type":   createInstanceToolType,
+				"source": "my-cloud-sql-source",
+			},
+		},
+	}
+}

tests/cloudsqlpg/cloud_sql_pg_integration_test.go

@@ -0,0 +1,298 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package cloudsqlpg
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"os"
+	"regexp"
+	"strings"
+	"testing"
+	"time"
+
+	"cloud.google.com/go/cloudsqlconn"
+	"github.com/google/uuid"
+	"github.com/googleapis/mcp-toolbox/internal/testutils"
+	"github.com/googleapis/mcp-toolbox/tests"
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+var (
+	CloudSQLPostgresSourceType = "cloud-sql-postgres"
+	CloudSQLPostgresToolType   = "postgres-sql"
+	CloudSQLPostgresProject    = os.Getenv("CLOUD_SQL_POSTGRES_PROJECT")
+	CloudSQLPostgresRegion     = os.Getenv("CLOUD_SQL_POSTGRES_REGION")
+	CloudSQLPostgresInstance   = os.Getenv("CLOUD_SQL_POSTGRES_INSTANCE")
+	CloudSQLPostgresDatabase   = os.Getenv("CLOUD_SQL_POSTGRES_DATABASE")
+	CloudSQLPostgresUser       = os.Getenv("CLOUD_SQL_POSTGRES_USER")
+	CloudSQLPostgresPass       = os.Getenv("CLOUD_SQL_POSTGRES_PASS")
+)
+
+func getCloudSQLPgVars(t *testing.T) map[string]any {
+	switch "" {
+	case CloudSQLPostgresProject:
+		t.Fatal("'CLOUD_SQL_POSTGRES_PROJECT' not set")
+	case CloudSQLPostgresRegion:
+		t.Fatal("'CLOUD_SQL_POSTGRES_REGION' not set")
+	case CloudSQLPostgresInstance:
+		t.Fatal("'CLOUD_SQL_POSTGRES_INSTANCE' not set")
+	case CloudSQLPostgresDatabase:
+		t.Fatal("'CLOUD_SQL_POSTGRES_DATABASE' not set")
+	case CloudSQLPostgresUser:
+		t.Fatal("'CLOUD_SQL_POSTGRES_USER' not set")
+	case CloudSQLPostgresPass:
+		t.Fatal("'CLOUD_SQL_POSTGRES_PASS' not set")
+	}
+
+	return map[string]any{
+		"type":     CloudSQLPostgresSourceType,
+		"project":  CloudSQLPostgresProject,
+		"instance": CloudSQLPostgresInstance,
+		"region":   CloudSQLPostgresRegion,
+		"database": CloudSQLPostgresDatabase,
+		"user":     CloudSQLPostgresUser,
+		"password": CloudSQLPostgresPass,
+	}
+}
+
+// Copied over from cloud_sql_pg.go
+func initCloudSQLPgConnectionPool(project, region, instance, ip_type, user, pass, dbname string) (*pgxpool.Pool, error) {
+	// Configure the driver to connect to the database
+	dsn := fmt.Sprintf("user=%s password=%s dbname=%s sslmode=disable", user, pass, dbname)
+	config, err := pgxpool.ParseConfig(dsn)
+	if err != nil {
+		return nil, fmt.Errorf("unable to parse connection uri: %w", err)
+	}
+
+	// Create a new dialer with options
+	dialOpts, err := tests.GetCloudSQLDialOpts(ip_type)
+	if err != nil {
+		return nil, err
+	}
+	d, err := cloudsqlconn.NewDialer(context.Background(), cloudsqlconn.WithDefaultDialOptions(dialOpts...))
+	if err != nil {
+		return nil, fmt.Errorf("unable to parse connection uri: %w", err)
+	}
+
+	// Tell the driver to use the Cloud SQL Go Connector to create connections
+	i := fmt.Sprintf("%s:%s:%s", project, region, instance)
+	config.ConnConfig.DialFunc = func(ctx context.Context, _ string, instance string) (net.Conn, error) {
+		return d.Dial(ctx, i)
+	}
+
+	// Interact with the driver directly as you normally would
+	pool, err := pgxpool.NewWithConfig(context.Background(), config)
+	if err != nil {
+		return nil, err
+	}
+	return pool, nil
+}
+
+func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) {
+	sourceConfig := getCloudSQLPgVars(t)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+
+	args := []string{"--enable-api"}
+
+	pool, err := initCloudSQLPgConnectionPool(CloudSQLPostgresProject, CloudSQLPostgresRegion, CloudSQLPostgresInstance, "public", CloudSQLPostgresUser, CloudSQLPostgresPass, CloudSQLPostgresDatabase)
+	if err != nil {
+		t.Fatalf("unable to create Cloud SQL connection pool: %s", err)
+	}
+
+	// Generate a unique ID
+	uniqueID := strings.ReplaceAll(uuid.New().String(), "-", "")
+
+	// This will execute after all tool tests complete (success, fail, or t.Fatal)
+	t.Cleanup(func() {
+		tests.CleanupPostgresTables(t, context.Background(), pool, uniqueID)
+	})
+
+	//Create table names using the UUID
+	tableNameParam := "param_table_" + uniqueID
+	tableNameAuth := "auth_table_" + uniqueID
+	tableNameTemplateParam := "template_param_table_" + uniqueID
+
+	// set up data for param tool
+	createParamTableStmt, insertParamTableStmt, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, paramTestParams := tests.GetPostgresSQLParamToolInfo(tableNameParam)
+	teardownTable1 := tests.SetupPostgresSQLTable(t, ctx, pool, createParamTableStmt, insertParamTableStmt, tableNameParam, paramTestParams)
+	defer teardownTable1(t)
+
+	// set up data for auth tool
+	createAuthTableStmt, insertAuthTableStmt, authToolStmt, authTestParams := tests.GetPostgresSQLAuthToolInfo(tableNameAuth)
+	teardownTable2 := tests.SetupPostgresSQLTable(t, ctx, pool, createAuthTableStmt, insertAuthTableStmt, tableNameAuth, authTestParams)
+	defer teardownTable2(t)
+
+	// Set up table for semantic search
+	vectorTableName, tearDownVectorTable := tests.SetupPostgresVectorTable(t, ctx, pool)
+	defer tearDownVectorTable(t)
+
+	// Write config into a file and pass it to command
+	toolsFile := tests.GetToolsConfig(sourceConfig, CloudSQLPostgresToolType, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, authToolStmt)
+	toolsFile = tests.AddExecuteSqlConfig(t, toolsFile, "postgres-execute-sql")
+	tmplSelectCombined, tmplSelectFilterCombined := tests.GetPostgresSQLTmplToolStatement()
+	toolsFile = tests.AddTemplateParamConfig(t, toolsFile, CloudSQLPostgresToolType, tmplSelectCombined, tmplSelectFilterCombined, "")
+
+	// Add semantic search tool config
+	insertStmt, searchStmt := tests.GetPostgresVectorSearchStmts(vectorTableName)
+	toolsFile = tests.AddSemanticSearchConfig(t, toolsFile, CloudSQLPostgresToolType, insertStmt, searchStmt)
+
+	toolsFile = tests.AddPostgresPrebuiltConfig(t, toolsFile)
+	cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
+	if err != nil {
+		t.Fatalf("command initialization returned an error: %s", err)
+	}
+	defer cleanup()
+
+	waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+	out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
+	if err != nil {
+		t.Logf("toolbox command logs: \n%s", out)
+		t.Fatalf("toolbox didn't start successfully: %s", err)
+	}
+
+	// Get configs for tests
+	select1Want, mcpMyFailToolWant, createTableStatement, mcpSelect1Want := tests.GetPostgresWants()
+
+	// Run tests
+	tests.RunToolGetTest(t)
+	tests.RunToolInvokeTest(t, select1Want)
+	tests.RunMCPToolCallMethod(t, mcpMyFailToolWant, mcpSelect1Want)
+	tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
+	tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam)
+
+	// Run Postgres prebuilt tool tests
+	tests.RunPostgresListTablesTest(t, tableNameParam, tableNameAuth, CloudSQLPostgresUser)
+	tests.RunPostgresListViewsTest(t, ctx, pool)
+	tests.RunPostgresListSchemasTest(t, ctx, pool, CloudSQLPostgresUser, uniqueID)
+	tests.RunPostgresListActiveQueriesTest(t, ctx, pool)
+	tests.RunPostgresListAvailableExtensionsTest(t)
+	tests.RunPostgresListInstalledExtensionsTest(t)
+	tests.RunPostgresDatabaseOverviewTest(t, ctx, pool)
+	tests.RunPostgresListTriggersTest(t, ctx, pool)
+	tests.RunPostgresListIndexesTest(t, ctx, pool)
+	tests.RunPostgresListSequencesTest(t, ctx, pool)
+	tests.RunPostgresListLocksTest(t, ctx, pool)
+	tests.RunPostgresReplicationStatsTest(t, ctx, pool)
+	tests.RunPostgresLongRunningTransactionsTest(t, ctx, pool)
+	tests.RunPostgresListQueryStatsTest(t, ctx, pool)
+	tests.RunPostgresGetColumnCardinalityTest(t, ctx, pool)
+	tests.RunPostgresListTableStatsTest(t, ctx, pool)
+	tests.RunPostgresListPublicationTablesTest(t, ctx, pool)
+	tests.RunPostgresListTableSpacesTest(t)
+	tests.RunPostgresListPgSettingsTest(t, ctx, pool)
+	tests.RunPostgresListDatabaseStatsTest(t, ctx, pool)
+	tests.RunPostgresListRolesTest(t, ctx, pool)
+	tests.RunPostgresListStoredProcedureTest(t, ctx, pool)
+	tests.RunSemanticSearchToolInvokeTest(t, "[]", "", "The quick brown fox")
+}
+
+// Test connection with different IP type
+func TestCloudSQLPgIpConnection(t *testing.T) {
+	sourceConfig := getCloudSQLPgVars(t)
+
+	tcs := []struct {
+		name   string
+		ipType string
+	}{
+		{
+			name:   "public ip",
+			ipType: "public",
+		},
+		{
+			name:   "private ip",
+			ipType: "private",
+		},
+	}
+	for _, tc := range tcs {
+		t.Run(tc.name, func(t *testing.T) {
+			sourceConfig["ipType"] = tc.ipType
+			err := tests.RunSourceConnectionTest(t, sourceConfig, CloudSQLPostgresToolType)
+			if err != nil {
+				t.Fatalf("Connection test failure: %s", err)
+			}
+		})
+	}
+}
+
+func TestCloudSQLPgIAMConnection(t *testing.T) {
+	getCloudSQLPgVars(t)
+	// service account email used for IAM should trim the suffix
+	serviceAccountEmail := strings.TrimSuffix(tests.ServiceAccountEmail, ".gserviceaccount.com")
+
+	noPassSourceConfig := map[string]any{
+		"type":     CloudSQLPostgresSourceType,
+		"project":  CloudSQLPostgresProject,
+		"instance": CloudSQLPostgresInstance,
+		"region":   CloudSQLPostgresRegion,
+		"database": CloudSQLPostgresDatabase,
+		"user":     serviceAccountEmail,
+	}
+
+	noUserSourceConfig := map[string]any{
+		"type":     CloudSQLPostgresSourceType,
+		"project":  CloudSQLPostgresProject,
+		"instance": CloudSQLPostgresInstance,
+		"region":   CloudSQLPostgresRegion,
+		"database": CloudSQLPostgresDatabase,
+		"password": "random",
+	}
+
+	noUserNoPassSourceConfig := map[string]any{
+		"type":     CloudSQLPostgresSourceType,
+		"project":  CloudSQLPostgresProject,
+		"instance": CloudSQLPostgresInstance,
+		"region":   CloudSQLPostgresRegion,
+		"database": CloudSQLPostgresDatabase,
+	}
+	tcs := []struct {
+		name         string
+		sourceConfig map[string]any
+		isErr        bool
+	}{
+		{
+			name:         "no user no pass",
+			sourceConfig: noUserNoPassSourceConfig,
+			isErr:        false,
+		},
+		{
+			name:         "no password",
+			sourceConfig: noPassSourceConfig,
+			isErr:        false,
+		},
+		{
+			name:         "no user",
+			sourceConfig: noUserSourceConfig,
+			isErr:        true,
+		},
+	}
+	for _, tc := range tcs {
+		t.Run(tc.name, func(t *testing.T) {
+			err := tests.RunSourceConnectionTest(t, tc.sourceConfig, CloudSQLPostgresToolType)
+			if err != nil {
+				if tc.isErr {
+					return
+				}
+				t.Fatalf("Connection test failure: %s", err)
+			}
+			if tc.isErr {
+				t.Fatalf("Expected error but test passed.")
+			}
+		})
+	}
+}

tests/cloudsqlpg/cloud_sql_pg_upgrade_precheck_test.go

@@ -0,0 +1,381 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package cloudsqlpg
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"regexp"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/googleapis/mcp-toolbox/internal/testutils"
+	"github.com/googleapis/mcp-toolbox/tests"
+	sqladmin "google.golang.org/api/sqladmin/v1"
+)
+
+var (
+	preCheckToolType = "postgres-upgrade-precheck"
+)
+
+type preCheckTransport struct {
+	transport http.RoundTripper
+	url       *url.URL
+}
+
+func (t *preCheckTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	if strings.HasPrefix(req.URL.String(), "https://sqladmin.googleapis.com") {
+		req.URL.Scheme = t.url.Scheme
+		req.URL.Host = t.url.Host
+	}
+	return t.transport.RoundTrip(req)
+}
+
+type preCheckHandler struct {
+	t            *testing.T
+	opCount      int
+	opResults    map[string][]*sqladmin.PreCheckResponse
+	opPollCounts map[string]int
+}
+
+func (h *preCheckHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	ua := r.Header.Get("User-Agent")
+	if !strings.Contains(ua, "genai-toolbox/") {
+		h.t.Errorf("User-Agent header not found in %q", ua)
+	}
+
+	if strings.Contains(r.URL.Path, "/operations/") {
+		h.handleOperations(w, r)
+		return
+	}
+
+	if strings.Contains(r.URL.Path, "/preCheckMajorVersionUpgrade") {
+		h.handlePreCheckV1(w, r)
+		return
+	}
+
+	http.Error(w, fmt.Sprintf("unhandled path: %s", r.URL.Path), http.StatusNotFound)
+}
+
+func (h *preCheckHandler) handlePreCheckV1(w http.ResponseWriter, r *http.Request) {
+	var body sqladmin.InstancesPreCheckMajorVersionUpgradeRequest
+	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+		h.t.Fatalf("failed to decode request body: %v", err)
+	}
+
+	if body.PreCheckMajorVersionUpgradeContext == nil || body.PreCheckMajorVersionUpgradeContext.TargetDatabaseVersion == "" {
+		http.Error(w, "missing targetDatabaseVersion", http.StatusBadRequest)
+		return
+	}
+
+	parts := strings.Split(r.URL.Path, "/")
+
+	if len(parts) < 7 {
+		msg := fmt.Sprintf("handlePreCheckV1: Expected 7 path parts, got %d for path %s", len(parts), r.URL.Path)
+		h.t.Errorf("%s", msg)
+		http.Error(w, msg, http.StatusBadRequest)
+		return
+	}
+
+	project := parts[3]
+	instanceName := parts[5]
+
+	h.opCount++
+	opName := fmt.Sprintf("op-%s-%s-%d", project, instanceName, h.opCount)
+
+	var preCheckResult []*sqladmin.PreCheckResponse
+	statusCode := http.StatusOK
+
+	switch instanceName {
+	case "instance-ok":
+		h.opResults[opName] = nil // This will make PreCheckResponse nil inside the context
+	case "instance-empty":
+		preCheckResult = []*sqladmin.PreCheckResponse{} // No issues
+		h.opResults[opName] = preCheckResult
+	case "instance-warnings":
+		preCheckResult = []*sqladmin.PreCheckResponse{
+			{
+				Message:         "This is a warning.",
+				MessageType:     "WARNING",
+				ActionsRequired: []string{"Check documentation."},
+			},
+		}
+		h.opResults[opName] = preCheckResult
+	case "instance-errors":
+		preCheckResult = []*sqladmin.PreCheckResponse{
+			{
+				Message:         "This is a critical error.",
+				MessageType:     "ERROR",
+				ActionsRequired: []string{"Fix this now."},
+			},
+		}
+		h.opResults[opName] = preCheckResult
+	case "instance-notfound":
+		http.Error(w, "Not authorized to access instance", http.StatusForbidden)
+		return
+	default:
+		msg := fmt.Sprintf("unhandled instance name in mock: %s", instanceName)
+		h.t.Errorf("handlePreCheckV1 default case: %s", msg)
+		http.Error(w, msg, http.StatusInternalServerError)
+		return
+	}
+
+	response := map[string]any{"name": opName, "status": "PENDING"}
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(statusCode)
+	if err := json.NewEncoder(w).Encode(response); err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+	}
+}
+
+func (h *preCheckHandler) handleOperations(w http.ResponseWriter, r *http.Request) {
+	parts := strings.Split(r.URL.Path, "/")
+	opName := parts[len(parts)-1]
+
+	h.opPollCounts[opName]++
+
+	result, ok := h.opResults[opName]
+	if !ok {
+		http.Error(w, fmt.Sprintf("operation not found: %s", opName), http.StatusNotFound)
+		return
+	}
+
+	status := "PENDING"
+	if h.opPollCounts[opName] > 1 {
+		status = "DONE"
+	}
+
+	opResponse := sqladmin.Operation{
+		Name:   opName,
+		Status: status,
+		Kind:   "sql#operation",
+	}
+
+	if status == "DONE" {
+		opResponse.PreCheckMajorVersionUpgradeContext = &sqladmin.PreCheckMajorVersionUpgradeContext{
+			PreCheckResponse: result, // This can be nil or empty
+		}
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	if err := json.NewEncoder(w).Encode(opResponse); err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+	}
+}
+
+// PreCheckResultItem holds the details of a single check result.
+type PreCheckResultItem struct {
+	Message         string   `json:"message"`
+	MessageType     string   `json:"messageType"`
+	ActionsRequired []string `json:"actionsRequired"`
+}
+
+// PreCheckAPIResponse holds the array of pre-check results.
+type PreCheckAPIResponse struct {
+	Items []PreCheckResultItem `json:"preCheckResponse"`
+}
+
+func TestPreCheckToolEndpoints(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
+	defer cancel()
+
+	handler := &preCheckHandler{
+		t:            t,
+		opResults:    make(map[string][]*sqladmin.PreCheckResponse),
+		opPollCounts: make(map[string]int),
+	}
+	server := httptest.NewServer(handler)
+	defer server.Close()
+
+	serverURL, err := url.Parse(server.URL)
+	if err != nil {
+		t.Fatalf("failed to parse server URL: %v", err)
+	}
+
+	originalTransport := http.DefaultClient.Transport
+	if originalTransport == nil {
+		originalTransport = http.DefaultTransport
+	}
+	http.DefaultClient.Transport = &preCheckTransport{
+		transport: originalTransport,
+		url:       serverURL,
+	}
+	t.Cleanup(func() {
+		http.DefaultClient.Transport = originalTransport
+	})
+
+	args := []string{"--enable-api"}
+	toolsFile := getPreCheckToolsConfig()
+	cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
+	if err != nil {
+		t.Fatalf("command initialization returned an error: %s", err)
+	}
+	defer cleanup()
+
+	waitCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
+	defer cancel()
+	_, err = testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
+	if err != nil {
+		t.Fatalf("toolbox didn't start successfully: %s", err)
+	}
+
+	tcs := []struct {
+		name        string
+		toolName    string
+		body        string
+		want        string
+		expectError bool
+		errorStatus int
+		errorMsg    string
+	}{
+		{
+			name:     "successful precheck - nil response in context",
+			toolName: "precheck-tool",
+			body:     `{"project": "p1", "instance": "instance-ok", "targetDatabaseVersion": "POSTGRES_18"}`,
+			want:     `{"preCheckResponse":[]}`, // Expect empty items list
+		},
+		{
+			name:     "successful precheck - empty issues",
+			toolName: "precheck-tool",
+			body:     `{"project": "p1", "instance": "instance-empty", "targetDatabaseVersion": "POSTGRES_18"}`,
+			want:     `{"preCheckResponse":[]}`,
+		},
+		{
+			name:     "successful precheck - with warnings",
+			toolName: "precheck-tool",
+			body:     `{"project": "p1", "instance": "instance-warnings", "targetDatabaseVersion": "POSTGRES_18"}`,
+			want:     `{"preCheckResponse":[{"actionsRequired":["Check documentation."],"type":"","message":"This is a warning.","messageType":"WARNING"}]}`,
+		},
+		{
+			name:     "successful precheck - with errors",
+			toolName: "precheck-tool",
+			body:     `{"project": "p1", "instance": "instance-errors", "targetDatabaseVersion": "POSTGRES_18"}`,
+			want:     `{"preCheckResponse":[{"actionsRequired":["Fix this now."],"type":"","message":"This is a critical error.","messageType":"ERROR"}]}`,
+		},
+		{
+			name:        "instance not found",
+			toolName:    "precheck-tool",
+			body:        `{"project": "p1", "instance": "instance-notfound", "targetDatabaseVersion": "POSTGRES_18"}`,
+			want:        `{"error":"failed to access GCP resource: googleapi: got HTTP response code 403 with body: Not authorized to access instance\n"}`,
+			expectError: true,
+			errorStatus: http.StatusInternalServerError,
+			errorMsg:    "failed to access GCP resource: googleapi: got HTTP response code 403",
+		},
+		{
+			name:     "missing required parameter - project",
+			toolName: "precheck-tool",
+			body:     `{"instance": "instance-ok", "targetDatabaseVersion": "POSTGRES_18"}`,
+			want:     `{"error":"parameter \"project\" is required"}`,
+		},
+		{
+			name:     "missing required parameter - instance",
+			toolName: "precheck-tool",
+			body:     `{"project": "p1", "targetDatabaseVersion": "POSTGRES_18"}`, // Missing instance
+			want:     `{"error":"parameter \"instance\" is required"}`,
+		},
+		{
+			name:     "missing parameter - targetDatabaseVersion",
+			toolName: "precheck-tool",
+			body:     `{"project": "p1", "instance": "instance-empty"}`, // Uses default POSTGRES_18
+			want:     `{"preCheckResponse":[]}`,
+		},
+	}
+
+	for _, tc := range tcs {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			api := fmt.Sprintf("http://127.0.0.1:5000/api/tool/%s/invoke", tc.toolName)
+			req, err := http.NewRequestWithContext(ctx, http.MethodPost, api, bytes.NewBufferString(tc.body))
+			if err != nil {
+				t.Fatalf("unable to create request: %s", err)
+			}
+			req.Header.Add("Content-type", "application/json")
+			req.Header.Add("Authorization", "Bearer FAKE_TOKEN")
+
+			resp, err := http.DefaultClient.Do(req)
+			if err != nil {
+				t.Fatalf("unable to send request: %s", err)
+			}
+			defer resp.Body.Close()
+
+			if tc.expectError {
+				if resp.StatusCode != tc.errorStatus {
+					bodyBytes, _ := io.ReadAll(resp.Body)
+					t.Fatalf("expected status %d but got %d: %s", tc.errorStatus, resp.StatusCode, string(bodyBytes))
+				}
+				bodyBytes, _ := io.ReadAll(resp.Body)
+				if !strings.Contains(string(bodyBytes), tc.errorMsg) {
+					t.Errorf("expected error message to contain %q, got %s", tc.errorMsg, string(bodyBytes))
+				}
+				return
+			}
+
+			if resp.StatusCode != http.StatusOK {
+				bodyBytes, _ := io.ReadAll(resp.Body)
+				t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes))
+			}
+
+			var result struct {
+				Result string `json:"result"`
+			}
+			if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
+				t.Fatalf("failed to decode response: %v", err)
+			}
+
+			var got PreCheckAPIResponse
+			if err := json.Unmarshal([]byte(result.Result), &got); err != nil {
+				t.Fatalf("failed to unmarshal result: %v", err)
+			}
+
+			var want PreCheckAPIResponse
+			if err := json.Unmarshal([]byte(tc.want), &want); err != nil {
+				t.Fatalf("failed to unmarshal want: %v", err)
+			}
+
+			if diff := cmp.Diff(want.Items, got.Items, cmp.Comparer(func(a, b PreCheckResultItem) bool {
+				return a.Message == b.Message && a.MessageType == b.MessageType && cmp.Equal(a.ActionsRequired, b.ActionsRequired)
+			})); diff != "" {
+				t.Errorf("unexpected result: diff (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
+
+func getPreCheckToolsConfig() map[string]any {
+	return map[string]any{
+		"sources": map[string]any{
+			"my-cloud-sql-source": map[string]any{
+				"type": "cloud-sql-admin",
+			},
+		},
+		"tools": map[string]any{
+			"precheck-tool": map[string]any{
+				"type":   preCheckToolType,
+				"source": "my-cloud-sql-source",
+				"authRequired": []string{
+					"https://www.googleapis.com/auth/cloud-platform",
+				},
+			},
+		},
+	}
+}

tests/cloudsqlpg/cloud_sql_pg_vectorassist_test.go

@@ -0,0 +1,583 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package cloudsqlpg
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"regexp"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/googleapis/mcp-toolbox/internal/testutils"
+	"github.com/googleapis/mcp-toolbox/tests"
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+func createPostgresExtension(t *testing.T, ctx context.Context, pool *pgxpool.Pool, extensionName string) func() {
+	createExtensionCmd := fmt.Sprintf("CREATE EXTENSION IF NOT EXISTS %s", extensionName)
+	_, err := pool.Exec(ctx, createExtensionCmd)
+	if err != nil {
+		t.Fatalf("failed to create extension: %v", err)
+	}
+	return func() {
+		dropExtensionCmd := fmt.Sprintf("DROP EXTENSION IF EXISTS %s", extensionName)
+		_, err := pool.Exec(ctx, dropExtensionCmd)
+		if err != nil {
+			t.Fatalf("failed to drop extension: %v", err)
+		}
+	}
+}
+
+// setupVectorAssistTable prepares the database extensions and test data needed
+// to test the definespec, modifyspec, applyspec, and generatequery tools.
+func setupVectorAssistTable(t *testing.T, ctx context.Context, pool *pgxpool.Pool) (string, func(t *testing.T), func()) {
+	// Install necessary extensions for VectorAssist
+	dropExtensionFunc := createPostgresExtension(t, ctx, pool, "vector_assist")
+
+	uniqueID := strings.ReplaceAll(uuid.New().String(), "-", "")
+	uniqueID = uniqueID[:10]
+	tableName := "vector_assist_test_" + uniqueID
+
+	// Create a table with vector data for defining/modifying/applying specs
+	createStmt := fmt.Sprintf(`
+		CREATE TABLE %s (
+			name TEXT,
+			category TEXT,
+			content TEXT,
+			embedding vector(3)
+		);
+	`, tableName)
+
+	_, err := pool.Exec(ctx, createStmt)
+	if err != nil {
+		t.Fatalf("failed to create vector assist test table: %v", err)
+	}
+
+	// Insert sample data to generate queries against
+	insertDataStmt := fmt.Sprintf(`
+		INSERT INTO %s (name, category, content, embedding)
+		VALUES
+		('Item 1', 'Document', 'Sample text document about AI', array_fill(0.1, ARRAY[3])::vector),
+		('Item 2', 'Document', 'Sample text document about databases', array_fill(0.2, ARRAY[3])::vector);
+	`, tableName)
+
+	_, err = pool.Exec(ctx, insertDataStmt)
+	if err != nil {
+		t.Fatalf("failed to insert data into vector assist table: %v", err)
+	}
+
+	// Return teardown function
+	teardown := func(t *testing.T) {
+		_, err := pool.Exec(context.Background(), fmt.Sprintf("DROP TABLE IF EXISTS %s;", tableName))
+		if err != nil {
+			t.Errorf("failed to drop vector assist table %s: %v", tableName, err)
+		}
+	}
+
+	return tableName, teardown, dropExtensionFunc
+}
+
+// TODO: Remove the test from this file and follow the existing test pattern
+// by calling the tests from cloudsqlpg_integration_test.go
+func TestVectorAssistIntegration(t *testing.T) {
+	// temporarily skip vector assist test since "Vector assist is temporarily
+	// disabled for all Cloud SQL for PostgreSQL instances"
+	// ref: https://docs.cloud.google.com/sql/docs/postgres/vector-assist-overview
+	t.Skip("Temporarily skipping vector assist test. Google had temporarily disabled vector assist for all CloudSQL for PostgreSQL instances.")
+
+	sourceConfig := getCloudSQLPgVars(t)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+
+	args := []string{"--enable-api"}
+
+	pool, err := initCloudSQLPgConnectionPool(CloudSQLPostgresProject, CloudSQLPostgresRegion, CloudSQLPostgresInstance, "public", CloudSQLPostgresUser, CloudSQLPostgresPass, CloudSQLPostgresDatabase)
+	if err != nil {
+		t.Fatalf("unable to create Cloud SQL connection pool: %s", err)
+	}
+
+	// Generate a unique ID
+	uniqueID := strings.ReplaceAll(uuid.New().String(), "-", "")
+
+	// This will execute after all tool tests complete (success, fail, or t.Fatal)
+	t.Cleanup(func() {
+		tests.CleanupPostgresTables(t, context.Background(), pool, uniqueID)
+	})
+
+	//Create table names using the UUID
+	tableNameParam := "param_table_" + uniqueID
+	tableNameAuth := "auth_table_" + uniqueID
+
+	// set up data for param tool
+	createParamTableStmt, insertParamTableStmt, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, paramTestParams := tests.GetPostgresSQLParamToolInfo(tableNameParam)
+	teardownTable1 := tests.SetupPostgresSQLTable(t, ctx, pool, createParamTableStmt, insertParamTableStmt, tableNameParam, paramTestParams)
+	defer teardownTable1(t)
+
+	// set up data for auth tool
+	createAuthTableStmt, insertAuthTableStmt, authToolStmt, authTestParams := tests.GetPostgresSQLAuthToolInfo(tableNameAuth)
+	teardownTable2 := tests.SetupPostgresSQLTable(t, ctx, pool, createAuthTableStmt, insertAuthTableStmt, tableNameAuth, authTestParams)
+	defer teardownTable2(t)
+
+	// // Set up data for vector assist tools
+	vectorAssistTableName, teardownVectorAssistTable, dropExtension := setupVectorAssistTable(t, ctx, pool)
+	defer teardownVectorAssistTable(t)
+	defer dropExtension()
+
+	// Write config into a file and pass it to command
+	toolsFile := tests.GetToolsConfig(sourceConfig, CloudSQLPostgresToolType, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, authToolStmt)
+	toolsFile = tests.AddExecuteSqlConfig(t, toolsFile, "postgres-execute-sql")
+	tmplSelectCombined, tmplSelectFilterCombined := tests.GetPostgresSQLTmplToolStatement()
+	toolsFile = tests.AddTemplateParamConfig(t, toolsFile, CloudSQLPostgresToolType, tmplSelectCombined, tmplSelectFilterCombined, "")
+
+	// Add vector assist tools to the configuration
+	toolsFile = AddVectorAssistConfig(t, toolsFile, "my-instance")
+
+	toolsFile = tests.AddPostgresPrebuiltConfig(t, toolsFile)
+	cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
+	if err != nil {
+		t.Fatalf("command initialization returned an error: %s", err)
+	}
+	defer cleanup()
+
+	waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+	out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
+	if err != nil {
+		t.Logf("toolbox command logs: \n%s", out)
+		t.Fatalf("toolbox didn't start successfully: %s", err)
+	}
+
+	// Run vectorassist tool tests
+	specID := "va_spec_001"
+	RunVectorAssistDefineSpecToolInvokeTest(t, ctx, pool, vectorAssistTableName, specID)
+	RunVectorAssistModifySpecToolInvokeTest(t, ctx, pool, specID)
+	RunVectorAssistApplySpecToolInvokeTest(t, ctx, pool, specID)
+	RunVectorAssistGenerateQueryToolInvokeTest(t, ctx, pool, specID)
+	RunVectorAssistImproveQueryRecallToolInvokeTest(t, ctx, pool, vectorAssistTableName)
+	RunVectorAssistListSpecsToolInvokeTest(t, ctx, pool, vectorAssistTableName)
+	RunVectorAssistGetSpecToolInvokeTest(t, ctx, pool, specID)
+	RunVectorAssistDeleteSpecToolInvokeTest(t, ctx, pool, specID)
+}
+
+// AddVectorAssistConfig appends the vector assist tool configurations to the given tools file.
+func AddVectorAssistConfig(t *testing.T, config map[string]any, sourceName string) map[string]any {
+	tools, ok := config["tools"].(map[string]any)
+	if !ok {
+		t.Fatalf("unable to get tools from config")
+	}
+
+	tools["define_spec"] = map[string]any{
+		"type":   "vector-assist-define-spec",
+		"source": sourceName,
+	}
+	tools["modify_spec"] = map[string]any{
+		"type":   "vector-assist-modify-spec",
+		"source": sourceName,
+	}
+	tools["apply_spec"] = map[string]any{
+		"type":   "vector-assist-apply-spec",
+		"source": sourceName,
+	}
+	tools["generate_query"] = map[string]any{
+		"type":   "vector-assist-generate-query",
+		"source": sourceName,
+	}
+	tools["improve_query_recall"] = map[string]any{
+		"type":   "vector-assist-improve-query-recall",
+		"source": sourceName,
+	}
+	tools["list_specs"] = map[string]any{
+		"type":   "vector-assist-list-specs",
+		"source": sourceName,
+	}
+	tools["get_spec"] = map[string]any{
+		"type":   "vector-assist-get-spec",
+		"source": sourceName,
+	}
+	tools["delete_spec"] = map[string]any{
+		"type":   "vector-assist-delete-spec",
+		"source": sourceName,
+	}
+	config["tools"] = tools
+	return config
+}
+
+func RunVectorAssistImproveQueryRecallToolInvokeTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool, tableName string) {
+	validPayload := fmt.Sprintf(`{
+    "schema_name": "public",
+		"table_name": "%s",
+    "vector_column_name": "embedding",
+    "index_name":  "%s_embedding_idx",
+    "top_k": 10,
+    "target_recall": 0.95,
+    "distance_func": "cosine"
+  }`,
+		tableName, tableName)
+
+	tcs := []struct {
+		name           string
+		requestBody    io.Reader
+		api            string
+		wantStatusCode int
+		wantContains   []string
+	}{
+		{
+			name:           "improve query recall valid parameters",
+			requestBody:    bytes.NewBuffer([]byte(validPayload)),
+			api:            "http://127.0.0.1:5000/api/tool/improve_query_recall/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`ef_search`,
+			},
+		},
+		{
+			name:           "improve query recall missing required table_name",
+			requestBody:    bytes.NewBuffer([]byte(`{"schema_name": "public"}`)),
+			api:            "http://127.0.0.1:5000/api/tool/improve_query_recall/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`"error"`,
+			},
+		},
+	}
+
+	for _, tc := range tcs {
+		runVectorAssistToolInvokeTest(t, tc)
+	}
+}
+
+func RunVectorAssistListSpecsToolInvokeTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool, tableName string) {
+	validPayload := fmt.Sprintf(`{
+    "table_name": "%s",
+    "column_name": "content"
+  }`,
+		tableName)
+
+	tcs := []struct {
+		name           string
+		requestBody    io.Reader
+		api            string
+		wantStatusCode int
+		wantContains   []string
+	}{
+		{
+			name:           "list specs valid parameters",
+			requestBody:    bytes.NewBuffer([]byte(validPayload)),
+			api:            "http://127.0.0.1:5000/api/tool/list_specs/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`va_spec_001`,
+			},
+		},
+		{
+			name:           "list specs missing required table_name",
+			requestBody:    bytes.NewBuffer([]byte(`{"column_name": "content"}`)),
+			api:            "http://127.0.0.1:5000/api/tool/list_specs/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`"error"`,
+			},
+		},
+	}
+
+	for _, tc := range tcs {
+		runVectorAssistToolInvokeTest(t, tc)
+	}
+}
+
+func RunVectorAssistGetSpecToolInvokeTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool, specID string) {
+	validPayload := fmt.Sprintf(`{
+    "spec_id": "%s"
+  }`,
+		specID)
+
+	tcs := []struct {
+		name           string
+		requestBody    io.Reader
+		api            string
+		wantStatusCode int
+		wantContains   []string
+	}{
+		{
+			name:           "get spec valid parameters",
+			requestBody:    bytes.NewBuffer([]byte(validPayload)),
+			api:            "http://127.0.0.1:5000/api/tool/get_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				specID,
+			},
+		},
+		{
+			name:           "get spec missing spec_id",
+			requestBody:    bytes.NewBuffer([]byte(`{}`)),
+			api:            "http://127.0.0.1:5000/api/tool/get_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`"error"`,
+			},
+		},
+	}
+
+	for _, tc := range tcs {
+		runVectorAssistToolInvokeTest(t, tc)
+	}
+}
+
+func RunVectorAssistDeleteSpecToolInvokeTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool, specID string) {
+	validPayload := fmt.Sprintf(`{
+    "spec_id": "%s"
+  }`,
+		specID)
+
+	tcs := []struct {
+		name           string
+		requestBody    io.Reader
+		api            string
+		wantStatusCode int
+		wantContains   []string
+	}{
+		{
+			name:           "delete spec valid parameters",
+			requestBody:    bytes.NewBuffer([]byte(validPayload)),
+			api:            "http://127.0.0.1:5000/api/tool/delete_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`delete_spec`,
+			},
+		},
+		{
+			name:           "delete spec missing spec_id",
+			requestBody:    bytes.NewBuffer([]byte(`{}`)),
+			api:            "http://127.0.0.1:5000/api/tool/delete_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`"error"`,
+			},
+		},
+	}
+
+	for _, tc := range tcs {
+		runVectorAssistToolInvokeTest(t, tc)
+	}
+}
+
+func RunVectorAssistDefineSpecToolInvokeTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool, tableName string, specID string) {
+	validPayload := fmt.Sprintf(`{
+		"table_name": "%s",
+		"schema_name": "public",
+		"spec_id": "%s",
+		"vector_column_name": "embedding",
+		"text_column_name": "content",
+		"vector_index_type": "hnsw",
+		"embeddings_available": true,
+		"num_vectors": 2,
+		"dimensionality": 3,
+		"embedding_model": "textembedding-gecko",
+		"prefilter_column_names": ["category"],
+		"distance_func": "cosine",
+		"quantization": "halfvec",
+		"memory_budget_kb": 1024,
+		"target_recall": 0.95,
+		"target_top_k": 10,
+		"tune_vector_index": true
+	}`, tableName, specID)
+
+	tcs := []struct {
+		name           string
+		requestBody    io.Reader
+		api            string
+		wantStatusCode int
+		wantContains   []string
+	}{
+		{
+			name:           "invoke define_spec with all valid parameters",
+			requestBody:    bytes.NewBuffer([]byte(validPayload)),
+			api:            "http://127.0.0.1:5000/api/tool/define_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			// Check for key identifiers instead of the entire JSON string
+			wantContains: []string{
+				`"vector_spec_id":"va_spec_001"`,
+				fmt.Sprintf(`"table_name":"%s"`, tableName),
+				`"recommendation_id"`, // Ensure a recommendation was generated
+			},
+		},
+		{
+			name:           "invoke define_spec with missing required table_name",
+			requestBody:    bytes.NewBuffer([]byte(`{"schema_name": "public", "spec_id": "va_spec_002"}`)),
+			api:            "http://127.0.0.1:5000/api/tool/define_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`"error"`,
+			},
+		},
+	}
+	for _, tc := range tcs {
+		runVectorAssistToolInvokeTest(t, tc)
+	}
+}
+
+func RunVectorAssistModifySpecToolInvokeTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool, specID string) {
+	validPayload := fmt.Sprintf(`{
+        "spec_id": "%s",
+        "memory_budget_kb": 2048,
+        "target_recall": 0.99
+    }`, specID)
+
+	tcs := []struct {
+		name           string
+		requestBody    io.Reader
+		api            string
+		wantStatusCode int
+		wantContains   []string
+	}{
+		{
+			name:           "modify existing spec with new constraints",
+			requestBody:    bytes.NewBuffer([]byte(validPayload)),
+			api:            "http://127.0.0.1:5000/api/tool/modify_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`"recommendation_id"`,
+			},
+		},
+		{
+			name:           "modify existing spec without required spec id",
+			requestBody:    bytes.NewBuffer([]byte(`{"target_recall": 0.99}`)),
+			api:            "http://127.0.0.1:5000/api/tool/modify_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`"error"`,
+			},
+		},
+	}
+	for _, tc := range tcs {
+		runVectorAssistToolInvokeTest(t, tc)
+	}
+}
+
+func RunVectorAssistApplySpecToolInvokeTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool, recommendationID string) {
+	validPayload := fmt.Sprintf(`{
+        "spec_id": "%s"
+    }`, recommendationID)
+
+	tcs := []struct {
+		name           string
+		requestBody    io.Reader
+		api            string
+		wantStatusCode int
+		wantContains   []string
+	}{
+		{
+			name:           "apply recommendation to database",
+			requestBody:    bytes.NewBuffer([]byte(validPayload)),
+			api:            "http://127.0.0.1:5000/api/tool/apply_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`{"apply_spec":true}`,
+			},
+		},
+		{
+			name:           "apply recommendation to database without spec id",
+			requestBody:    bytes.NewBuffer([]byte(`{"schema_name": "public"}`)),
+			api:            "http://127.0.0.1:5000/api/tool/apply_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`"error"`,
+			},
+		},
+	}
+
+	for _, tc := range tcs {
+		runVectorAssistToolInvokeTest(t, tc)
+	}
+}
+
+func RunVectorAssistGenerateQueryToolInvokeTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool, specID string) {
+	validPayload := fmt.Sprintf(`{
+        "spec_id": "%s",
+        "search_text": "What is the capital of France?",
+        "top_k": 5
+    }`, specID)
+
+	tcs := []struct {
+		name           string
+		requestBody    io.Reader
+		api            string
+		wantStatusCode int
+		wantContains   []string
+	}{
+		{
+			name:           "generate SQL for vector search",
+			requestBody:    bytes.NewBuffer([]byte(validPayload)),
+			api:            "http://127.0.0.1:5000/api/tool/generate_query/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`"generate_query"`,
+				`LIMIT 5`,
+			},
+		},
+	}
+
+	for _, tc := range tcs {
+		runVectorAssistToolInvokeTest(t, tc)
+	}
+}
+
+func runVectorAssistToolInvokeTest(t *testing.T, tc struct {
+	name           string
+	requestBody    io.Reader
+	api            string
+	wantStatusCode int
+	wantContains   []string
+}) {
+	t.Run(tc.name, func(t *testing.T) {
+		resp, body := tests.RunRequest(t, http.MethodPost, tc.api, tc.requestBody, nil)
+
+		if resp.StatusCode != tc.wantStatusCode {
+			t.Fatalf("tool %s: wrong status code: got %d, want %d, body: %s", tc.api, resp.StatusCode, tc.wantStatusCode, string(body))
+		}
+
+		if tc.wantStatusCode != http.StatusOK {
+			return
+		}
+
+		// Unmarshal the standard response wrapper
+		var bodyWrapper struct {
+			Result json.RawMessage `json:"result"`
+		}
+		if err := json.Unmarshal(body, &bodyWrapper); err != nil {
+			t.Fatalf("error decoding response wrapper: %v", err)
+		}
+
+		// Handle the double-unmarshal logic for stringified results
+		var resultString string
+		if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
+			resultString = string(bodyWrapper.Result)
+		}
+
+		// Verification loop
+		for _, expectedSubstr := range tc.wantContains {
+			if !strings.Contains(resultString, expectedSubstr) {
+				t.Errorf("Expected result to contain %q, but it did not.\nFull result: %s", expectedSubstr, resultString)
+			}
+		}
+	})
+}