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

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 @@
 import (
 	"context"
 	"fmt"
-	"net/url"
-	"os"
 	"regexp"
 	"strings"
 	"testing"
@@ -28,37 +26,23 @@
 	"github.com/googleapis/genai-toolbox/internal/testutils"
 	"github.com/googleapis/genai-toolbox/tests"
 	"github.com/jackc/pgx/v5/pgxpool"
+	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 +51,8 @@
 	}
 }
 
-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 +61,36 @@
 }
 
 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")
+	if err != nil {
+		t.Fatalf("failed to start container: %s", err)
+	}
+	defer tccockroachdbContainer.Terminate(ctx)
+
+	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, err := tccockroachdbContainer.ConnectionString(ctx)
+	if err != nil {
+		t.Fatalf("failed to get connection string: %s", err)
+	}
+
+	connString += "?sslmode=disable"
+
+	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)
 	}