test(source/cockroachdb): migrate integration tests to testcontainers

.ci/integration.cloudbuild.yaml

@@ -359,10 +359,8 @@ steps:
     entrypoint: /bin/bash
     env:
       - "GOPATH=/gopath"
-      - "COCKROACHDB_DATABASE=$_DATABASE_NAME"
-      - "COCKROACHDB_PORT=$_COCKROACHDB_PORT"
       - "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
-    secretEnv: ["COCKROACHDB_USER", "COCKROACHDB_HOST","CLIENT_ID"]
+    secretEnv: ["CLIENT_ID"]
     volumes:
       - name: "go"
         path: "/gopath"
@@ -1173,11 +1171,6 @@ availableSecrets:
       env: MARIADB_HOST
     - versionName: projects/$PROJECT_ID/secrets/mongodb_uri/versions/latest
       env: MONGODB_URI
-    - versionName: projects/$PROJECT_ID/secrets/cockroachdb_user/versions/latest
-      env: COCKROACHDB_USER
-    - versionName: projects/$PROJECT_ID/secrets/cockroachdb_host/versions/latest
-      env: COCKROACHDB_HOST
-
 
 options:
   logging: CLOUD_LOGGING_ONLY

go.mod

@@ -49,6 +49,7 @@ require (
 	github.com/snowflakedb/gosnowflake v1.18.1
 	github.com/spf13/cobra v1.10.1
 	github.com/testcontainers/testcontainers-go v0.40.0
+	github.com/testcontainers/testcontainers-go/modules/cockroachdb v0.40.0
 	github.com/thlib/go-timezone-local v0.0.7
 	github.com/trinodb/trino-go-client v0.330.0
 	github.com/valkey-io/valkey-go v1.0.68

go.sum

@@ -1357,6 +1357,8 @@ github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu
 github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
 github.com/testcontainers/testcontainers-go v0.40.0 h1:pSdJYLOVgLE8YdUY2FHQ1Fxu+aMnb6JfVz1mxk7OeMU=
 github.com/testcontainers/testcontainers-go v0.40.0/go.mod h1:FSXV5KQtX2HAMlm7U3APNyLkkap35zNLxukw9oBi/MY=
+github.com/testcontainers/testcontainers-go/modules/cockroachdb v0.40.0 h1:UNYfrnFV9mkO93Sw6hqRA5KbE9DsAvDeYKD4GDConiE=
+github.com/testcontainers/testcontainers-go/modules/cockroachdb v0.40.0/go.mod h1:O8By1J/1y726YYk7obTIXxfv2OzonVe+ORq9Z+K+fDg=
 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/pretty v1.0.0/go.mod h1:XNkn88O1ChpSDQmQeStsy+sBenx6DDtFZJxhVysOjyk=

tests/cockroachdb/cockroachdb_integration_test.go

@@ -17,8 +17,6 @@ package cockroachdb
 import (
 	"context"
 	"fmt"
-	"net/url"
-	"os"
 	"regexp"
 	"strings"
 	"testing"
@@ -28,37 +26,24 @@ import (
 	"github.com/googleapis/genai-toolbox/internal/testutils"
 	"github.com/googleapis/genai-toolbox/tests"
 	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/testcontainers/testcontainers-go"
+	tccockroachdb "github.com/testcontainers/testcontainers-go/modules/cockroachdb"
 )
 
 var (
 	CockroachDBSourceType = "cockroachdb"
 	CockroachDBToolType   = "cockroachdb-sql"
-	CockroachDBDatabase   = getEnvOrDefault("COCKROACHDB_DATABASE", "defaultdb")
-	CockroachDBHost       = getEnvOrDefault("COCKROACHDB_HOST", "localhost")
-	CockroachDBPort       = getEnvOrDefault("COCKROACHDB_PORT", "26257")
-	CockroachDBUser       = getEnvOrDefault("COCKROACHDB_USER", "root")
-	CockroachDBPass       = getEnvOrDefault("COCKROACHDB_PASS", "")
+	CockroachDBDatabase   = "defaultdb"
 )
 
-func getEnvOrDefault(key, defaultValue string) string {
-	if value := os.Getenv(key); value != "" {
-		return value
-	}
-	return defaultValue
-}
-
-func getCockroachDBVars(t *testing.T) map[string]any {
-	if CockroachDBHost == "" {
-		t.Skip("COCKROACHDB_HOST not set, skipping CockroachDB integration test")
-	}
-
+func getCockroachDBVars(host, port string) map[string]any {
 	return map[string]any{
 		"type":           CockroachDBSourceType,
-		"host":           CockroachDBHost,
-		"port":           CockroachDBPort,
+		"host":           host,
+		"port":           port,
 		"database":       CockroachDBDatabase,
-		"user":           CockroachDBUser,
-		"password":       CockroachDBPass,
+		"user":           "root",
+		"password":       "",
 		"maxRetries":     5,
 		"retryBaseDelay": "500ms",
 		"queryParams": map[string]string{
@@ -67,15 +52,8 @@ func getCockroachDBVars(t *testing.T) map[string]any {
 	}
 }
 
-func initCockroachDBConnectionPool(host, port, user, pass, dbname string) (*pgxpool.Pool, error) {
-	connURL := &url.URL{
-		Scheme:   "postgres",
-		User:     url.UserPassword(user, pass),
-		Host:     fmt.Sprintf("%s:%s", host, port),
-		Path:     dbname,
-		RawQuery: "sslmode=disable&application_name=cockroachdb-integration-test",
-	}
-	pool, err := pgxpool.New(context.Background(), connURL.String())
+func initCockroachDBConnectionPool(connString string) (*pgxpool.Pool, error) {
+	pool, err := pgxpool.New(context.Background(), connString)
 	if err != nil {
 		return nil, fmt.Errorf("unable to create connection pool: %w", err)
 	}
@@ -84,13 +62,37 @@ func initCockroachDBConnectionPool(host, port, user, pass, dbname string) (*pgxp
 }
 
 func TestCockroachDB(t *testing.T) {
-	sourceConfig := getCockroachDBVars(t)
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancel()
 
+	tccockroachdbContainer, err := tccockroachdb.Run(ctx, "cockroachdb/cockroach:latest-v23.1",
+		testcontainers.WithCmd("start-single-node", "--insecure"),
+	)
+	if err != nil {
+		t.Fatalf("failed to start container: %s", err)
+	}
+	t.Cleanup(func() {
+		if err := tccockroachdbContainer.Terminate(context.Background()); err != nil {
+			t.Logf("failed to terminate testcontainer: %s", err)
+		}
+	})
+
+	host, err := tccockroachdbContainer.Host(ctx)
+	if err != nil {
+		t.Fatalf("failed to get host: %s", err)
+	}
+
+	port, err := tccockroachdbContainer.MappedPort(ctx, "26257/tcp")
+	if err != nil {
+		t.Fatalf("failed to get port: %s", err)
+	}
+
+	connString := fmt.Sprintf("postgres://root@%s:%s/defaultdb?sslmode=disable", host, port.Port())
+
+	sourceConfig := getCockroachDBVars(host, port.Port())
 	var args []string
 
-	pool, err := initCockroachDBConnectionPool(CockroachDBHost, CockroachDBPort, CockroachDBUser, CockroachDBPass, CockroachDBDatabase)
+	pool, err := initCockroachDBConnectionPool(connString)
 	if err != nil {
 		t.Fatalf("unable to create cockroachdb connection pool: %s", err)
 	}
@@ -112,11 +114,6 @@ func TestCockroachDB(t *testing.T) {
 	// 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