diff --git a/.ci/integration.cloudbuild.yaml b/.ci/integration.cloudbuild.yaml index 519c218301be..be34659b32bb 100644 --- a/.ci/integration.cloudbuild.yaml +++ b/.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 diff --git a/go.mod b/go.mod index f9e9ad96fbf2..669fd24963ee 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index dc0c0ebb2f70..641157ecdb38 100644 --- a/go.sum +++ b/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= diff --git a/tests/cockroachdb/cockroachdb_integration_test.go b/tests/cockroachdb/cockroachdb_integration_test.go index 464b82238e41..a83bc3aded1f 100644 --- a/tests/cockroachdb/cockroachdb_integration_test.go +++ b/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