Skip to content

Commit 01d4bfb

Browse files
anubhav756anubhav756
committed
test(source/alloydbomni): create MCP integration tests
1 parent 4fe514a commit 01d4bfb

File tree

3 files changed

+1725
-53
lines changed

3 files changed

+1725
-53
lines changed

‎tests/alloydbomni/alloydb_omni_integration_test.go‎

Lines changed: 0 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -28,14 +28,6 @@ import (
2828
"github.com/googleapis/genai-toolbox/internal/testutils"
2929
"github.com/googleapis/genai-toolbox/tests"
3030
"github.com/jackc/pgx/v5/pgxpool"
31-
"github.com/testcontainers/testcontainers-go"
32-
"github.com/testcontainers/testcontainers-go/wait"
33-
)
34-
35-
var (
36-
AlloyDBUser = "postgres"
37-
AlloyDBPass = "mysecretpassword"
38-
AlloyDBDatabase = "postgres"
3931
)
4032

4133
func buildPostgresURL(host, port, user, pass, dbname string) *url.URL {
@@ -57,51 +49,6 @@ func initPostgresConnectionPool(host, port, user, pass, dbname string) (*pgxpool
5749
return pool, nil
5850
}
5951

60-
func setupAlloyDBContainer(ctx context.Context, t *testing.T) (string, string, func()) {
61-
t.Helper()
62-
63-
req := testcontainers.ContainerRequest{
64-
Image: "google/alloydbomni:16.9.0-ubi9", // Pinning version for stability
65-
ExposedPorts: []string{"5432/tcp"},
66-
Env: map[string]string{
67-
"POSTGRES_PASSWORD": AlloyDBPass,
68-
},
69-
WaitingFor: wait.ForAll(
70-
wait.ForLog("database system was shut down at"),
71-
wait.ForLog("database system is ready to accept connections"),
72-
wait.ForExposedPort(),
73-
),
74-
}
75-
76-
container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
77-
ContainerRequest: req,
78-
Started: true,
79-
})
80-
if err != nil {
81-
t.Fatalf("failed to start alloydb container: %s", err)
82-
}
83-
84-
cleanup := func() {
85-
if err := container.Terminate(ctx); err != nil {
86-
t.Fatalf("failed to terminate container: %s", err)
87-
}
88-
}
89-
90-
host, err := container.Host(ctx)
91-
if err != nil {
92-
cleanup()
93-
t.Fatalf("failed to get container host: %s", err)
94-
}
95-
96-
mappedPort, err := container.MappedPort(ctx, "5432")
97-
if err != nil {
98-
cleanup()
99-
t.Fatalf("failed to get container mapped port: %s", err)
100-
}
101-
102-
return host, mappedPort.Port(), cleanup
103-
}
104-
10552
func TestAlloyDBOmni(t *testing.T) {
10653
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
10754
defer cancel()

‎tests/alloydbomni/alloydb_omni_mcp_test.go‎

Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
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 alloydbomni
16+
17+
import (
18+
"context"
19+
"os"
20+
"regexp"
21+
"strings"
22+
"testing"
23+
"time"
24+
25+
"github.com/google/uuid"
26+
"github.com/googleapis/genai-toolbox/internal/testutils"
27+
"github.com/googleapis/genai-toolbox/tests"
28+
"github.com/testcontainers/testcontainers-go"
29+
"github.com/testcontainers/testcontainers-go/wait"
30+
)
31+
32+
var (
33+
AlloyDBUser = "postgres"
34+
AlloyDBPass = "mysecretpassword"
35+
AlloyDBDatabase = "postgres"
36+
)
37+
38+
func setupAlloyDBContainer(ctx context.Context, t *testing.T) (string, string, func()) {
39+
t.Helper()
40+
41+
req := testcontainers.ContainerRequest{
42+
Image: "google/alloydbomni:16.9.0-ubi9", // Pinning version for stability
43+
ExposedPorts: []string{"5432/tcp"},
44+
Env: map[string]string{
45+
"POSTGRES_PASSWORD": AlloyDBPass,
46+
},
47+
WaitingFor: wait.ForAll(
48+
wait.ForLog("database system was shut down at"),
49+
wait.ForLog("database system is ready to accept connections"),
50+
wait.ForExposedPort(),
51+
),
52+
}
53+
54+
container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
55+
ContainerRequest: req,
56+
Started: true,
57+
})
58+
if err != nil {
59+
t.Fatalf("failed to start alloydb container: %s", err)
60+
}
61+
62+
cleanup := func() {
63+
if err := container.Terminate(ctx); err != nil {
64+
t.Fatalf("failed to terminate container: %s", err)
65+
}
66+
}
67+
68+
host, err := container.Host(ctx)
69+
if err != nil {
70+
cleanup()
71+
t.Fatalf("failed to get container host: %s", err)
72+
}
73+
74+
mappedPort, err := container.MappedPort(ctx, "5432")
75+
if err != nil {
76+
cleanup()
77+
t.Fatalf("failed to get container mapped port: %s", err)
78+
}
79+
80+
return host, mappedPort.Port(), cleanup
81+
}
82+
83+
func TestAlloyDBOmniMCP(t *testing.T) {
84+
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
85+
defer cancel()
86+
87+
AlloyDBHost, AlloyDBPort, containerCleanup := setupAlloyDBContainer(ctx, t)
88+
defer containerCleanup()
89+
90+
os.Setenv("ALLOYDB_OMNI_HOST", AlloyDBHost)
91+
os.Setenv("ALLOYDB_OMNI_PORT", AlloyDBPort)
92+
os.Setenv("ALLOYDB_OMNI_USER", AlloyDBUser)
93+
os.Setenv("ALLOYDB_OMNI_PASSWORD", AlloyDBPass)
94+
os.Setenv("ALLOYDB_OMNI_DATABASE", AlloyDBDatabase)
95+
96+
// Generate a unique ID
97+
uniqueID := strings.ReplaceAll(uuid.New().String(), "-", "")
98+
99+
args := []string{"--prebuilt", "alloydb-omni"}
100+
101+
pool, err := tests.InitPostgresConnectionPool(AlloyDBHost, AlloyDBPort, AlloyDBUser, AlloyDBPass, AlloyDBDatabase)
102+
if err != nil {
103+
t.Fatalf("unable to create alloydb connection pool: %s", err)
104+
}
105+
106+
cmd, cleanup, err := tests.StartCmd(ctx, map[string]any{}, args...)
107+
if err != nil {
108+
t.Fatalf("command initialization returned an error: %s", err)
109+
}
110+
defer cleanup()
111+
112+
// Wait for server to be ready
113+
waitCtx, waitCancel := context.WithTimeout(ctx, 30*time.Second)
114+
defer waitCancel()
115+
116+
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
117+
if err != nil {
118+
t.Logf("toolbox command logs: \n%s", out)
119+
t.Fatalf("toolbox didn't start successfully: %s", err)
120+
}
121+
122+
t.Logf("AlloyDB Omni container started on %s:%s", AlloyDBHost, AlloyDBPort)
123+
t.Logf("Toolbox server started with output: %s", out)
124+
125+
tests.RunMCPPostgresListViewsTest(t, ctx, pool)
126+
tests.RunMCPPostgresListSchemasTest(t, ctx, pool, AlloyDBUser, uniqueID)
127+
tests.RunMCPPostgresListActiveQueriesTest(t, ctx, pool)
128+
tests.RunMCPPostgresListAvailableExtensionsTest(t)
129+
tests.RunMCPPostgresListInstalledExtensionsTest(t)
130+
tests.RunMCPPostgresListTriggersTest(t, ctx, pool)
131+
tests.RunMCPPostgresListSequencesTest(t, ctx, pool)
132+
tests.RunMCPPostgresListIndexesTest(t, ctx, pool)
133+
tests.RunMCPPostgresListStoredProcedureTest(t, ctx, pool)
134+
tests.RunMCPPostgresDatabaseOverviewTest(t, ctx, pool)
135+
tests.RunMCPPostgresListLocksTest(t, ctx, pool)
136+
tests.RunMCPPostgresLongRunningTransactionsTest(t, ctx, pool)
137+
tests.RunMCPPostgresReplicationStatsTest(t, ctx, pool)
138+
tests.RunMCPPostgresGetColumnCardinalityTest(t, ctx, pool)
139+
tests.RunMCPPostgresListTableStatsTest(t, ctx, pool)
140+
tests.RunMCPPostgresListPublicationTablesTest(t, ctx, pool)
141+
tests.RunMCPPostgresListTableSpacesTest(t)
142+
tests.RunMCPPostgresListPgSettingsTest(t, ctx, pool)
143+
tests.RunMCPPostgresListDatabaseStatsTest(t, ctx, pool)
144+
tests.RunMCPPostgresListRolesTest(t, ctx, pool)
145+
146+
toolsToTest := map[string]string{
147+
"list_autovacuum_configurations": `{}`,
148+
"list_memory_configurations": `{}`,
149+
"list_top_bloated_tables": `{"limit": 10}`,
150+
"list_replication_slots": `{}`,
151+
"list_invalid_indexes": `{}`,
152+
"get_query_plan": `{"query": "SELECT 1"}`,
153+
"list_columnar_configurations": `{}`,
154+
"list_columnar_recommended_columns": `{}`,
155+
}
156+
tests.RunMCPStatementToolsTest(t, toolsToTest)
157+
}

0 commit comments

Comments
 (0)