veza/veza-backend-api/internal/testutils/setup.go

package testutils

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/testcontainers/testcontainers-go"
	"github.com/testcontainers/testcontainers-go/modules/postgres"
	"github.com/testcontainers/testcontainers-go/wait"
	"go.uber.org/zap"
)

var (
	pgContainer   *postgres.PostgresContainer
	pgDSN         string
	containerOnce sync.Once
	pgErr         error
)

// GetTestContainerDB ensures the postgres container is running and returns the DSN.
// It uses a singleton pattern to start the container only once per test run.
// If DATABASE_URL is set (e.g. CI with real PostgreSQL), uses it instead of testcontainer.
func GetTestContainerDB(ctx context.Context) (string, error) {
	containerOnce.Do(func() {
		if dsn := os.Getenv("DATABASE_URL"); dsn != "" {
			pgDSN = dsn
			pgErr = nil
			return
		}
		pgErr = setupPostgresContainer(ctx)
	})
	return pgDSN, pgErr
}

func setupPostgresContainer(ctx context.Context) error {
	// Find project root relative to this file
	// This file is in internal/testutils/setup.go
	_, filename, _, _ := runtime.Caller(0)
	projectRoot := filepath.Join(filepath.Dir(filename), "../..")
	migrationsDir := filepath.Join(projectRoot, "migrations")

	// Collect migration files
	files, err := os.ReadDir(migrationsDir)
	if err != nil {
		return fmt.Errorf("failed to read migrations dir: %w", err)
	}

	var migrationFiles []string
	for _, f := range files {
		// MOD-P1-001: Exclude cleanup migrations that may fail if tables don't exist yet
		// These migrations are meant to be run on existing databases, not fresh ones
		if strings.HasSuffix(f.Name(), ".sql") && !strings.Contains(f.Name(), "000000_cleanup") {
			migrationFiles = append(migrationFiles, filepath.Join(migrationsDir, f.Name()))
		}
	}
	sort.Strings(migrationFiles) // Ensure alphabetical order (001_, 002_, ...)

	// MOD-P1-001: Retry container startup with exponential backoff
	// Use a simple logger for testcontainers (zap.L() may not be initialized in tests)
	logger := zap.NewNop()
	if zap.L() != nil {
		logger = zap.L()
	}

	var containerErr error
	maxRetries := 3
	retryDelay := 2 * time.Second

	for attempt := 1; attempt <= maxRetries; attempt++ {
		logger.Info("Starting PostgreSQL testcontainer",
			zap.Int("attempt", attempt),
			zap.Int("max_retries", maxRetries),
			zap.Int("migration_files", len(migrationFiles)),
		)

		// Start Postgres container with improved wait strategy
		pgContainer, containerErr = postgres.Run(ctx,
			"postgres:15-alpine",
			postgres.WithDatabase("veza_test"),
			postgres.WithUsername("veza"),
			postgres.WithPassword("veza"),
			postgres.WithInitScripts(migrationFiles...),
			testcontainers.WithWaitStrategy(
				wait.ForLog("database system is ready to accept connections").
					WithOccurrence(2).
					WithStartupTimeout(90*time.Second)), // Increased timeout from 60s to 90s
		)

		if containerErr == nil {
			logger.Info("PostgreSQL testcontainer started successfully",
				zap.Int("attempt", attempt),
			)
			break // Success
		}

		// Log retry attempt
		logger.Warn("Failed to start PostgreSQL testcontainer, retrying",
			zap.Int("attempt", attempt),
			zap.Int("max_retries", maxRetries),
			zap.Error(containerErr),
		)

		if attempt < maxRetries {
			backoff := retryDelay * time.Duration(attempt) // Exponential backoff
			logger.Info("Waiting before retry",
				zap.Duration("backoff", backoff),
			)
			time.Sleep(backoff)
		}
	}

	if containerErr != nil {
		logger.Error("Failed to start PostgreSQL testcontainer after all retries",
			zap.Int("max_retries", maxRetries),
			zap.Error(containerErr),
		)
		return fmt.Errorf("failed to start postgres container after %d attempts: %w", maxRetries, containerErr)
	}

	var dsnErr error
	pgDSN, dsnErr = pgContainer.ConnectionString(ctx, "sslmode=disable")
	if dsnErr != nil {
		return fmt.Errorf("failed to get connection string: %w", dsnErr)
	}

	return nil
}

// TerminateContainer allows manual termination if needed (mostly for cleanup)
func TerminateContainer(ctx context.Context) error {
	if pgContainer != nil {
		return pgContainer.Terminate(ctx)
	}
	return nil
}