Skip to content

test(source/cloudsqlpg): create MCP integration tests #2994

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
anubhav756 wants to merge 6 commits into
base: anubhav-mcp-alloydbpg
Choose a base branch
from
Open

test(source/cloudsqlpg): create MCP integration tests #2994

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
454e8aa
test(source/cloudsqlpg): create MCP integration tests
anubhav756 Apr 8, 2026
5f6fb9e
fix: Fix nil 2
anubhav756 Apr 8, 2026
d667484
chore: Remove comments
anubhav756 Apr 8, 2026
5a32837
fix: compliance [0] for alloydbainl
anubhav756 Apr 8, 2026
a10c605
fix: [0] issue for common harness
anubhav756 Apr 8, 2026
aa1537c
fix: context passing
anubhav756 Apr 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 1 addition & 4 deletions tests/alloydbainl/alloydb_ai_nl_mcp_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -198,10 +198,7 @@ func TestAlloyDBAINLCallTool(t *testing.T) {
if mcpResp.Result.IsError {
t.Fatalf("expected success result, got tool error: %v", mcpResp.Result)
}
if len(mcpResp.Result.Content) == 0 {
t.Fatalf("expected at least one content item, got none")
}
got := mcpResp.Result.Content[0].Text
got := tests.GetMCPResultText(mcpResp)
if got != tc.want {
t.Fatalf("unexpected value: got %q, want %q", got, tc.want)
}
Expand Down
78 changes: 1 addition & 77 deletions tests/cloudsqlpg/cloud_sql_pg_integration_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,100 +16,24 @@ package cloudsqlpg

import (
"context"
"fmt"
"net"
"os"
"regexp"
"strings"
"testing"
"time"

"cloud.google.com/go/cloudsqlconn"
"github.com/google/uuid"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-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)
pool, err := initCloudSQLPgConnectionPool(ctx, CloudSQLPostgresProject, CloudSQLPostgresRegion, CloudSQLPostgresInstance, "public", CloudSQLPostgresUser, CloudSQLPostgresPass, CloudSQLPostgresDatabase)
if err != nil {
t.Fatalf("unable to create Cloud SQL connection pool: %s", err)
}
Expand Down
194 changes: 194 additions & 0 deletions tests/cloudsqlpg/cloud_sql_pg_mcp_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,194 @@
// 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 (
"context"
"fmt"
"net"
"os"
"regexp"
"strings"
"testing"
"time"

"cloud.google.com/go/cloudsqlconn"
"github.com/google/uuid"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-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")
}
Comment on lines +46 to +59
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot Apr 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This switch "" pattern for checking environment variables is clever, but it has a downside: it fails the test on the first missing variable. This can be frustrating for developers who then have to fix one variable and rerun tests, only to find another one is missing.

A more developer-friendly approach is to check for all required variables and report all missing ones at once. This gives a complete picture of the required setup.

	vars := map[string]string{
		"CLOUD_SQL_POSTGRES_PROJECT":  CloudSQLPostgresProject,
		"CLOUD_SQL_POSTGRES_REGION":   CloudSQLPostgresRegion,
		"CLOUD_SQL_POSTGRES_INSTANCE": CloudSQLPostgresInstance,
		"CLOUD_SQL_POSTGRES_DATABASE": CloudSQLPostgresDatabase,
		"CLOUD_SQL_POSTGRES_USER":     CloudSQLPostgresUser,
		"CLOUD_SQL_POSTGRES_PASS":     CloudSQLPostgresPass,
	}
	var missing []string
	for key, val := range vars {
		if val == "" {
			missing = append(missing, key)
		}
	}
	if len(missing) > 0 {
		t.Fatalf("Missing required environment variables: %s", strings.Join(missing, ", "))
	}


return map[string]any{
"type": CloudSQLPostgresSourceType,
"project": CloudSQLPostgresProject,
"instance": CloudSQLPostgresInstance,
"region": CloudSQLPostgresRegion,
"database": CloudSQLPostgresDatabase,
"user": CloudSQLPostgresUser,
"password": CloudSQLPostgresPass,
}
}

func initCloudSQLPgConnectionPool(ctx context.Context, project, region, instance, ip_type, user, pass, dbname string) (*pgxpool.Pool, error) {
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(ctx, 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(ctx, config)
if err != nil {
return nil, err
}
return pool, nil
}

func TestCloudSQLPgListTools(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
defer cancel()

args := []string{"--prebuilt", "cloud-sql-postgres"}

cmd, cleanup, err := tests.StartCmd(ctx, map[string]any{}, args...)
if err != nil {
t.Fatalf("command initialization returned an error: %v", err)
}
defer cleanup()

waitCtx, cancelWait := context.WithTimeout(ctx, 10*time.Second)
defer cancelWait()
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: %v", err)
}

// We expect standard Postgres tools to be listed
_, tools, err := tests.GetMCPToolsList(t, nil)
if err != nil {
t.Fatalf("failed to get tools list: %v", err)
}

if len(tools) == 0 {
t.Errorf("expected tools to be listed, got none")
}
}

func TestCloudSQLPgCallTool(t *testing.T) {
getCloudSQLPgVars(t)
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot Apr 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The getCloudSQLPgVars function is called here, but its return value is ignored. This is confusing because the function name suggests it gets something, but it's being used only for the side-effect of validating environment variables.

To improve clarity, consider renaming the function to requireCloudSQLPgVars to make its purpose as a validator explicit.

Also, TestCloudSQLPgListTools would benefit from this check at the beginning, as it also depends on these environment variables to start the server with the prebuilt configuration.

ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
defer cancel()

pool, err := initCloudSQLPgConnectionPool(ctx, CloudSQLPostgresProject, CloudSQLPostgresRegion, CloudSQLPostgresInstance, "public", CloudSQLPostgresUser, CloudSQLPostgresPass, CloudSQLPostgresDatabase)
if err != nil {
t.Fatalf("unable to create Cloud SQL connection pool: %s", err)
}
defer pool.Close()

uniqueID := strings.ReplaceAll(uuid.New().String(), "-", "")

args := []string{"--prebuilt", "cloud-sql-postgres"}

cmd, cleanup, err := tests.StartCmd(ctx, map[string]any{}, args...)
if err != nil {
t.Fatalf("command initialization returned an error: %v", err)
}
defer cleanup()

waitCtx, cancelWait := context.WithTimeout(ctx, 10*time.Second)
defer cancelWait()
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: %v", err)
}

// Run shared Postgres tests
tests.RunMCPPostgresListViewsTest(t, ctx, pool)
tests.RunMCPPostgresListSchemasTest(t, ctx, pool, CloudSQLPostgresUser, uniqueID)
tests.RunMCPPostgresListActiveQueriesTest(t, ctx, pool)
tests.RunMCPPostgresListAvailableExtensionsTest(t)
tests.RunMCPPostgresListInstalledExtensionsTest(t)
tests.RunMCPPostgresDatabaseOverviewTest(t, ctx, pool)
tests.RunMCPPostgresListTriggersTest(t, ctx, pool)
tests.RunMCPPostgresListIndexesTest(t, ctx, pool)
tests.RunMCPPostgresListSequencesTest(t, ctx, pool)
tests.RunMCPPostgresLongRunningTransactionsTest(t, ctx, pool)
tests.RunMCPPostgresListLocksTest(t, ctx, pool)
tests.RunMCPPostgresReplicationStatsTest(t, ctx, pool)
tests.RunMCPPostgresGetColumnCardinalityTest(t, ctx, pool)
tests.RunMCPPostgresListTableStatsTest(t, ctx, pool)
tests.RunMCPPostgresListPublicationTablesTest(t, ctx, pool)
tests.RunMCPPostgresListTableSpacesTest(t)
tests.RunMCPPostgresListPgSettingsTest(t, ctx, pool)
tests.RunMCPPostgresListDatabaseStatsTest(t, ctx, pool)
tests.RunMCPPostgresListRolesTest(t, ctx, pool)
tests.RunMCPPostgresListStoredProcedureTest(t, ctx, pool)

toolsToTest := map[string]string{
"list_autovacuum_configurations": `{}`,
"list_memory_configurations": `{}`,
"list_top_bloated_tables": `{"limit": 10}`,
"list_replication_slots": `{}`,
"list_invalid_indexes": `{}`,
"get_query_plan": `{"query": "SELECT 1"}`,
}
tests.RunMCPStatementToolsTest(t, toolsToTest)
}
Loading
Loading