From 0da8a74c71c4a4fb6a8832b3492738a42a5d27bc Mon Sep 17 00:00:00 2001 From: duwenxin99 Date: Thu, 26 Mar 2026 14:43:08 -0400 Subject: [PATCH 1/3] test(couchbase): refactor integration tests to use testcontainers --- tests/couchbase/couchbase_integration_test.go | 136 +++++++++--------- 1 file changed, 68 insertions(+), 68 deletions(-) diff --git a/tests/couchbase/couchbase_integration_test.go b/tests/couchbase/couchbase_integration_test.go index 1d7bb6bfc499..3488e117e13b 100644 --- a/tests/couchbase/couchbase_integration_test.go +++ b/tests/couchbase/couchbase_integration_test.go @@ -17,7 +17,6 @@ package couchbase import ( "context" "fmt" - "os" "regexp" "strings" "testing" @@ -27,48 +26,30 @@ import ( "github.com/google/uuid" "github.com/googleapis/genai-toolbox/internal/testutils" "github.com/googleapis/genai-toolbox/tests" + tccouchbase "github.com/testcontainers/testcontainers-go/modules/couchbase" ) const ( couchbaseSourceType = "couchbase" couchbaseToolType = "couchbase-sql" + defaultBucketName = "test-bucket" + defaultUser = "Administrator" + defaultPass = "password" ) -var ( - couchbaseConnection = os.Getenv("COUCHBASE_CONNECTION") - couchbaseBucket = os.Getenv("COUCHBASE_BUCKET") - couchbaseScope = os.Getenv("COUCHBASE_SCOPE") - couchbaseUser = os.Getenv("COUCHBASE_USER") - couchbasePass = os.Getenv("COUCHBASE_PASS") -) - -// getCouchbaseVars validates and returns Couchbase configuration variables -func getCouchbaseVars(t *testing.T) map[string]any { - switch "" { - case couchbaseConnection: - t.Fatal("'COUCHBASE_CONNECTION' not set") - case couchbaseBucket: - t.Fatal("'COUCHBASE_BUCKET' not set") - case couchbaseScope: - t.Fatal("'COUCHBASE_SCOPE' not set") - case couchbaseUser: - t.Fatal("'COUCHBASE_USER' not set") - case couchbasePass: - t.Fatal("'COUCHBASE_PASS' not set") - } - +// getCouchbaseVars generates config using dynamic container info +func getCouchbaseVars(connectionString string) map[string]any { return map[string]any{ "type": couchbaseSourceType, - "connectionString": couchbaseConnection, - "bucket": couchbaseBucket, - "scope": couchbaseScope, - "username": couchbaseUser, - "password": couchbasePass, + "connectionString": connectionString, + "bucket": defaultBucketName, + "scope": "_default", // Testcontainers default + "username": defaultUser, + "password": defaultPass, "queryScanConsistency": 2, } } -// initCouchbaseCluster initializes a connection to the Couchbase cluster func initCouchbaseCluster(connectionString, username, password string) (*gocb.Cluster, error) { opts := gocb.ClusterOptions{ Authenticator: gocb.PasswordAuthenticator{ @@ -76,7 +57,6 @@ func initCouchbaseCluster(connectionString, username, password string) (*gocb.Cl Password: password, }, } - cluster, err := gocb.Connect(connectionString, opts) if err != nil { return nil, fmt.Errorf("gocb.Connect: %w", err) @@ -85,73 +65,93 @@ func initCouchbaseCluster(connectionString, username, password string) (*gocb.Cl } func TestCouchbaseToolEndpoints(t *testing.T) { - sourceConfig := getCouchbaseVars(t) - ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute) defer cancel() - var args []string + // Start Couchbase Container + cbContainer, err := tccouchbase.Run(ctx, "couchbase/server:7.2.0", + tccouchbase.WithBucket(tccouchbase.NewBucket(defaultBucketName)), + tccouchbase.WithCredentials(defaultUser, defaultPass), + ) + if err != nil { + t.Fatalf("failed to start couchbase container: %s", err) + } + t.Cleanup(func() { + if err := cbContainer.Terminate(context.Background()); err != nil { + t.Logf("failed to terminate container: %s", err) + } + }) + + connectionString, err := cbContainer.ConnectionString(ctx) + if err != nil { + t.Fatalf("failed to get connection string: %s", err) + } - cluster, err := initCouchbaseCluster(couchbaseConnection, couchbaseUser, couchbasePass) + // 2. Initialize SDK Client for Setup + cluster, err := initCouchbaseCluster(connectionString, defaultUser, defaultPass) if err != nil { t.Fatalf("unable to create Couchbase connection: %s", err) } defer cluster.Close(nil) - // Create collection names with UUID + sourceConfig := getCouchbaseVars(connectionString) + scopeName := "_default" + + // Prepare Data collectionNameParam := "param_" + strings.ReplaceAll(uuid.New().String(), "-", "") collectionNameAuth := "auth_" + strings.ReplaceAll(uuid.New().String(), "-", "") collectionNameTemplateParam := "template_param_" + strings.ReplaceAll(uuid.New().String(), "-", "") - // Set up data for param tool - paramToolStatement, idParamToolStmt, nameParamToolStmt, arrayToolStatement, paramTestParams := getCouchbaseParamToolInfo(collectionNameParam) - teardownCollection1 := setupCouchbaseCollection(t, ctx, cluster, couchbaseBucket, couchbaseScope, collectionNameParam, paramTestParams) - defer teardownCollection1(t) + paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, paramTestParams := getCouchbaseParamToolInfo(collectionNameParam) + teardown1 := setupCouchbaseCollection(t, ctx, cluster, defaultBucketName, scopeName, collectionNameParam, paramTestParams) + defer teardown1(t) - // Set up data for auth tool - authToolStatement, authTestParams := getCouchbaseAuthToolInfo(collectionNameAuth) - teardownCollection2 := setupCouchbaseCollection(t, ctx, cluster, couchbaseBucket, couchbaseScope, collectionNameAuth, authTestParams) - defer teardownCollection2(t) + authToolStmt, authTestParams := getCouchbaseAuthToolInfo(collectionNameAuth) + teardown2 := setupCouchbaseCollection(t, ctx, cluster, defaultBucketName, scopeName, collectionNameAuth, authTestParams) + defer teardown2(t) - // Setup up table for template param tool tmplSelectCombined, tmplSelectFilterCombined, tmplSelectAll, params3 := getCouchbaseTemplateParamToolInfo() - teardownCollection3 := setupCouchbaseCollection(t, ctx, cluster, couchbaseBucket, couchbaseScope, collectionNameTemplateParam, params3) - defer teardownCollection3(t) + teardown3 := setupCouchbaseCollection(t, ctx, cluster, defaultBucketName, scopeName, collectionNameTemplateParam, params3) + defer teardown3(t) - // Write config into a file and pass it to command - toolsFile := tests.GetToolsConfig(sourceConfig, couchbaseToolType, paramToolStatement, idParamToolStmt, nameParamToolStmt, arrayToolStatement, authToolStatement) + // 4. Configure Toolbox + toolsFile := tests.GetToolsConfig(sourceConfig, couchbaseToolType, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, authToolStmt) toolsFile = tests.AddTemplateParamConfig(t, toolsFile, couchbaseToolType, tmplSelectCombined, tmplSelectFilterCombined, tmplSelectAll) - cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...) + cmd, cleanup, err := tests.StartCmd(ctx, toolsFile) if err != nil { - t.Fatalf("command initialization returned an error: %s", err) + t.Fatalf("command initialization failed: %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) + // Wait for server + waitCtx, waitCancel := context.WithTimeout(ctx, 20*time.Second) + defer waitCancel() + if _, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out); err != nil { + t.Fatalf("toolbox didn't start: %s", err) } - // Get configs for tests + // Assertions select1Want := "[{\"$1\":1}]" mcpMyFailToolWant := `{"jsonrpc":"2.0","id":"invoke-fail-tool","result":{"content":[{"type":"text","text":"error processing request: unable to execute query: parsing failure | {\"statement\":\"SELEC 1;\"` mcpSelect1Want := `{"jsonrpc":"2.0","id":"invoke my-auth-required-tool","result":{"content":[{"type":"text","text":"{\"$1\":1}"}]}}` tmplSelectId1Want := "[{\"age\":21,\"id\":1,\"name\":\"Alex\"}]" selectAllWant := "[{\"age\":21,\"id\":1,\"name\":\"Alex\"},{\"age\":100,\"id\":2,\"name\":\"Alice\"}]" - // Run tests - tests.RunToolGetTest(t) - tests.RunToolInvokeTest(t, select1Want) - tests.RunMCPToolCallMethod(t, mcpMyFailToolWant, mcpSelect1Want) - tests.RunToolInvokeWithTemplateParameters(t, collectionNameTemplateParam, - tests.WithTmplSelectId1Want(tmplSelectId1Want), - tests.WithSelectAllWant(selectAllWant), - tests.DisableDdlTest(), - tests.DisableInsertTest(), - ) + t.Run("GeneralTests", func(t *testing.T) { + tests.RunToolGetTest(t) + tests.RunToolInvokeTest(t, select1Want) + tests.RunMCPToolCallMethod(t, mcpMyFailToolWant, mcpSelect1Want) + }) + + t.Run("TemplateTests", func(t *testing.T) { + tests.RunToolInvokeWithTemplateParameters(t, collectionNameTemplateParam, + tests.WithTmplSelectId1Want(tmplSelectId1Want), + tests.WithSelectAllWant(selectAllWant), + tests.DisableDdlTest(), + tests.DisableInsertTest(), + ) + }) } // setupCouchbaseCollection creates a scope and collection and inserts test data From 9020a12b835d8753800a1a5a8751a90cd1229dab Mon Sep 17 00:00:00 2001 From: duwenxin99 Date: Thu, 26 Mar 2026 14:48:41 -0400 Subject: [PATCH 2/3] update integration test yaml --- .ci/integration.cloudbuild.yaml | 12 +----------- tests/couchbase/couchbase_integration_test.go | 5 +++-- 2 files changed, 4 insertions(+), 13 deletions(-) diff --git a/.ci/integration.cloudbuild.yaml b/.ci/integration.cloudbuild.yaml index 7c21e803a343..6bf5a335c9cd 100644 --- a/.ci/integration.cloudbuild.yaml +++ b/.ci/integration.cloudbuild.yaml @@ -756,11 +756,9 @@ steps: entrypoint: /bin/bash env: - "GOPATH=/gopath" - - "COUCHBASE_SCOPE=$_COUCHBASE_SCOPE" - - "COUCHBASE_BUCKET=$_COUCHBASE_BUCKET" - "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL" secretEnv: - ["COUCHBASE_CONNECTION", "COUCHBASE_USER", "COUCHBASE_PASS", "CLIENT_ID"] + ["CLIENT_ID"] volumes: - name: "go" path: "/gopath" @@ -1432,12 +1430,6 @@ availableSecrets: env: MSSQL_USER - versionName: projects/$PROJECT_ID/secrets/mssql_pass/versions/latest env: MSSQL_PASS - - versionName: projects/$PROJECT_ID/secrets/couchbase_connection/versions/latest - env: COUCHBASE_CONNECTION - - versionName: projects/$PROJECT_ID/secrets/couchbase_user/versions/latest - env: COUCHBASE_USER - - versionName: projects/$PROJECT_ID/secrets/couchbase_pass/versions/latest - env: COUCHBASE_PASS - versionName: projects/$PROJECT_ID/secrets/memorystore_redis_address/versions/latest env: REDIS_ADDRESS - versionName: projects/$PROJECT_ID/secrets/memorystore_redis_pass/versions/latest @@ -1549,8 +1541,6 @@ substitutions: _MSSQL_HOST: 127.0.0.1 _MSSQL_PORT: "1433" _DGRAPHURL: "https://play.dgraph.io" - _COUCHBASE_BUCKET: "couchbase-bucket" - _COUCHBASE_SCOPE: "couchbase-scope" _LOOKER_LOCATION: "us" _LOOKER_PROJECT: "149671255749" _LOOKER_VERIFY_SSL: "true" diff --git a/tests/couchbase/couchbase_integration_test.go b/tests/couchbase/couchbase_integration_test.go index 3488e117e13b..02266f5c6bd1 100644 --- a/tests/couchbase/couchbase_integration_test.go +++ b/tests/couchbase/couchbase_integration_test.go @@ -50,6 +50,7 @@ func getCouchbaseVars(connectionString string) map[string]any { } } +// initCouchbaseCluster initializes a connection to the Couchbase cluster func initCouchbaseCluster(connectionString, username, password string) (*gocb.Cluster, error) { opts := gocb.ClusterOptions{ Authenticator: gocb.PasswordAuthenticator{ @@ -87,7 +88,7 @@ func TestCouchbaseToolEndpoints(t *testing.T) { t.Fatalf("failed to get connection string: %s", err) } - // 2. Initialize SDK Client for Setup + // Set up Clouchbase cluster cluster, err := initCouchbaseCluster(connectionString, defaultUser, defaultPass) if err != nil { t.Fatalf("unable to create Couchbase connection: %s", err) @@ -114,7 +115,7 @@ func TestCouchbaseToolEndpoints(t *testing.T) { teardown3 := setupCouchbaseCollection(t, ctx, cluster, defaultBucketName, scopeName, collectionNameTemplateParam, params3) defer teardown3(t) - // 4. Configure Toolbox + // Configure Toolbox toolsFile := tests.GetToolsConfig(sourceConfig, couchbaseToolType, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, authToolStmt) toolsFile = tests.AddTemplateParamConfig(t, toolsFile, couchbaseToolType, tmplSelectCombined, tmplSelectFilterCombined, tmplSelectAll) From d6f6ead7d82ef6f530df1aacaf1e45b72ed2ef11 Mon Sep 17 00:00:00 2001 From: duwenxin99 Date: Thu, 26 Mar 2026 14:52:01 -0400 Subject: [PATCH 3/3] fix lint --- go.mod | 4 ++++ go.sum | 8 ++++++++ tests/couchbase/couchbase_integration_test.go | 4 ++-- 3 files changed, 14 insertions(+), 2 deletions(-) diff --git a/go.mod b/go.mod index 26ae4a652153..ffe624e9e9f7 100644 --- a/go.mod +++ b/go.mod @@ -54,6 +54,7 @@ require ( github.com/spf13/pflag v1.0.10 github.com/testcontainers/testcontainers-go v0.41.0 github.com/testcontainers/testcontainers-go/modules/cockroachdb v0.41.0 + github.com/testcontainers/testcontainers-go/modules/couchbase v0.41.0 github.com/thlib/go-timezone-local v0.0.7 github.com/trinodb/trino-go-client v0.333.0 github.com/valkey-io/valkey-go v1.0.73 @@ -239,6 +240,9 @@ require ( github.com/sirupsen/logrus v1.9.4 // indirect github.com/spiffe/go-spiffe/v2 v2.6.0 // indirect github.com/stretchr/testify v1.11.1 // indirect + github.com/tidwall/gjson v1.17.1 // indirect + github.com/tidwall/match v1.1.1 // indirect + github.com/tidwall/pretty v1.2.0 // indirect github.com/tklauser/go-sysconf v0.3.16 // indirect github.com/tklauser/numcpus v0.11.0 // indirect github.com/xdg-go/pbkdf2 v1.0.0 // indirect diff --git a/go.sum b/go.sum index f59f628d384d..2c6f0642b270 100644 --- a/go.sum +++ b/go.sum @@ -576,9 +576,17 @@ github.com/testcontainers/testcontainers-go v0.41.0 h1:mfpsD0D36YgkxGj2LrIyxuwQ9 github.com/testcontainers/testcontainers-go v0.41.0/go.mod h1:pdFrEIfaPl24zmBjerWTTYaY0M6UHsqA1YSvsoU40MI= github.com/testcontainers/testcontainers-go/modules/cockroachdb v0.41.0 h1:OnTmWiIbYcWP/TsE+RE6xBuFPY6Ouc4tQFQTqBbCQFA= github.com/testcontainers/testcontainers-go/modules/cockroachdb v0.41.0/go.mod h1:G+4O/i3AlRZto+u3/VQNZ3EQtyAQhUypuZjNBQB5gjI= +github.com/testcontainers/testcontainers-go/modules/couchbase v0.41.0 h1:8Lv3ZnvTzxHRHeD1TWyiFZg+jkN1t3TJqnfoUu0eX2c= +github.com/testcontainers/testcontainers-go/modules/couchbase v0.41.0/go.mod h1:lkGzggZ3Sp1MYzVI26r0vWwQdeaSrby2wLD8HM3TiW0= github.com/thlib/go-timezone-local v0.0.7 h1:fX8zd3aJydqLlTs/TrROrIIdztzsdFV23OzOQx31jII= github.com/thlib/go-timezone-local v0.0.7/go.mod h1:/Tnicc6m/lsJE0irFMA0LfIwTBo4QP7A8IfyIv4zZKI= +github.com/tidwall/gjson v1.17.1 h1:wlYEnwqAHgzmhNUFfw7Xalt2JzQvsMx2Se4PcoFCT/U= +github.com/tidwall/gjson v1.17.1/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= +github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA= +github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= github.com/tidwall/pretty v1.0.0/go.mod h1:XNkn88O1ChpSDQmQeStsy+sBenx6DDtFZJxhVysOjyk= +github.com/tidwall/pretty v1.2.0 h1:RWIZEg2iJ8/g6fDDYzMpobmaoGh5OLl4AXtGUGPcqCs= +github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= github.com/tklauser/go-sysconf v0.3.16 h1:frioLaCQSsF5Cy1jgRBrzr6t502KIIwQ0MArYICU0nA= github.com/tklauser/go-sysconf v0.3.16/go.mod h1:/qNL9xxDhc7tx3HSRsLWNnuzbVfh3e7gh/BmM179nYI= github.com/tklauser/numcpus v0.11.0 h1:nSTwhKH5e1dMNsCdVBukSZrURJRoHbSEQjdEbY+9RXw= diff --git a/tests/couchbase/couchbase_integration_test.go b/tests/couchbase/couchbase_integration_test.go index 02266f5c6bd1..61f7e9a09f6f 100644 --- a/tests/couchbase/couchbase_integration_test.go +++ b/tests/couchbase/couchbase_integration_test.go @@ -71,8 +71,8 @@ func TestCouchbaseToolEndpoints(t *testing.T) { // Start Couchbase Container cbContainer, err := tccouchbase.Run(ctx, "couchbase/server:7.2.0", - tccouchbase.WithBucket(tccouchbase.NewBucket(defaultBucketName)), - tccouchbase.WithCredentials(defaultUser, defaultPass), + tccouchbase.WithAdminCredentials(defaultUser, defaultPass), + tccouchbase.WithBuckets(tccouchbase.NewBucket(defaultBucketName)), ) if err != nil { t.Fatalf("failed to start couchbase container: %s", err)