gochat/backend/internal/repository/testdb_helper.go

package repository

import (
	"fmt"
	"os"
	"sync"
	"testing"

	"gorm.io/driver/postgres"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"

	"github.com/gochat/gochat/internal/canned"
	"github.com/gochat/gochat/internal/csat"
	"github.com/gochat/gochat/internal/model"
)

// actualDB tracks whether we actually connected to PG or fell back to SQLite.
// This is set once during the first setupTestDB call and used by skipIfSQLite.
var (
	actualDBIsPG    bool
	actualDBChecked bool
	actualDBCheckMu sync.Mutex
	pgSchemaMu      sync.Mutex
	pgMigrated      = make(map[string]bool)
	pgExtensionsSet bool
	sqliteSchemaMu  sync.Mutex
	sqliteDB        *gorm.DB
	sqliteMigrated  = make(map[string]bool)
)

// wantPostgres returns true when GOCHAT_TEST_DB != "sqlite".
func wantPostgres() bool {
	return os.Getenv("GOCHAT_TEST_DB") != "sqlite"
}

// isActuallyPG returns true if the test database is actually PostgreSQL
// (not a SQLite fallback). Probes PG once and caches the result.
func isActuallyPG() bool {
	actualDBCheckMu.Lock()
	defer actualDBCheckMu.Unlock()

	if actualDBChecked {
		return actualDBIsPG
	}

	if !wantPostgres() {
		actualDBIsPG = false
		actualDBChecked = true
		return false
	}

	// Probe PG connectivity
	dsn := pgDSN()
	db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	if err != nil {
		actualDBIsPG = false
		actualDBChecked = true
		return false
	}
	sqlDB, _ := db.DB()
	sqlDB.Close()
	actualDBIsPG = true
	actualDBChecked = true
	return true
}

// pgDSN returns the PostgreSQL connection string for testing.
// Priority: GOCHAT_TEST_DB_URL env > default localhost DSN.
func pgDSN() string {
	dsn := os.Getenv("GOCHAT_TEST_DB_URL")
	if dsn != "" {
		return dsn
	}
	return "host=localhost port=5432 user=postgres password=postgres dbname=gochat_test sslmode=disable"
}

// setupTestDB creates an isolated test database.
// When PG is available, connects to PostgreSQL; otherwise uses SQLite in-memory.
// The returned DB has all GoChat models auto-migrated.
func setupTestDB(t *testing.T, models ...interface{}) *gorm.DB {
	t.Helper()

	allModels := append(defaultTestModels(), models...)

	if isActuallyPG() {
		db, err := gorm.Open(postgres.Open(pgDSN()), &gorm.Config{
			DisableForeignKeyConstraintWhenMigrating: true,
			Logger:                                   logger.Default.LogMode(logger.Silent),
		})
		if err != nil {
			t.Fatalf("failed to connect to PostgreSQL test db: %v\nSet GOCHAT_TEST_DB=sqlite to use SQLite instead", err)
		}
		if err := migratePGModels(db, allModels); err != nil {
			t.Fatalf("failed to prepare PostgreSQL test schema: %v", err)
		}
		tx := db.Begin()
		if tx.Error != nil {
			t.Fatalf("failed to begin PostgreSQL test transaction: %v", tx.Error)
		}

		t.Cleanup(func() {
			_ = tx.Rollback().Error
			sqlDB, _ := db.DB()
			sqlDB.Close()
		})

		return tx
	}

	return openSQLiteDB(t, allModels)
}

// setupBenchmarkDB creates an isolated database for benchmarks.
func setupBenchmarkDB(b *testing.B, models ...interface{}) *gorm.DB {
	b.Helper()

	allModels := append(defaultTestModels(), models...)

	if isActuallyPG() {
		db, err := gorm.Open(postgres.Open(pgDSN()), &gorm.Config{
			DisableForeignKeyConstraintWhenMigrating: true,
			Logger:                                   logger.Default.LogMode(logger.Silent),
		})
		if err != nil {
			b.Fatalf("failed to connect to PostgreSQL benchmark db: %v", err)
		}
		if err := migratePGModels(db, allModels); err != nil {
			b.Fatalf("failed to prepare PostgreSQL benchmark schema: %v", err)
		}
		tx := db.Begin()
		if tx.Error != nil {
			b.Fatalf("failed to begin PostgreSQL benchmark transaction: %v", tx.Error)
		}

		b.Cleanup(func() {
			_ = tx.Rollback().Error
			sqlDB, _ := db.DB()
			sqlDB.Close()
		})

		return tx
	}

	return openSQLiteDBBench(b, allModels)
}

func migratePGModels(db *gorm.DB, models []interface{}) error {
	pgSchemaMu.Lock()
	defer pgSchemaMu.Unlock()

	if !pgExtensionsSet {
		if err := db.Exec("CREATE EXTENSION IF NOT EXISTS vector; CREATE EXTENSION IF NOT EXISTS pg_trgm").Error; err != nil {
			return fmt.Errorf("enable extensions: %w", err)
		}
		pgExtensionsSet = true
	}

	pending := make([]interface{}, 0, len(models))
	keys := make([]string, 0, len(models))
	for _, m := range models {
		key := fmt.Sprintf("%T", m)
		if pgMigrated[key] {
			continue
		}
		pgMigrated[key] = true
		pending = append(pending, m)
		keys = append(keys, key)
	}
	if len(pending) == 0 {
		return nil
	}
	if err := db.AutoMigrate(pending...); err != nil {
		for _, key := range keys {
			delete(pgMigrated, key)
		}
		return fmt.Errorf("auto-migrate: %w", err)
	}
	return nil
}

// defaultTestModels returns the standard set of models needed for most tests.
func defaultTestModels() []interface{} {
	return []interface{}{
		&model.Account{},
		&model.User{},
		&model.AccountUser{},
		&model.Inbox{},
		&model.Contact{},
		&model.ContactInbox{},
		&model.Conversation{},
		&model.Message{},
		&model.Notification{},
		&model.NotificationPreference{},
		&model.Team{},
		&model.TeamMember{},
		&canned.CannedResponse{},
		&csat.CsatSurveyResponse{},
		&model.Tag{},
		&model.ConversationLabel{},
		&model.ContactLabel{},
		&model.ContactExport{},
		&model.DataImport{},
		&model.PlatformApp{},
		&model.Permissible{},
		&model.AccessToken{},
		&model.ConversationParticipant{},
		&model.DraftMessage{},
		&model.Audit{},
		&model.CustomRole{},
		&model.AgentCapacityPolicy{},
		&model.InboxCapacityLimit{},
		&model.Company{},
		&model.CompanyNote{},
		&model.Mention{},
		&model.IntegrationHook{},
		&model.IntegrationApp{},
		&model.BackgroundJob{},
	}
}

func openSQLiteDB(t *testing.T, models []interface{}) *gorm.DB {
	t.Helper()
	sqliteSchemaMu.Lock()
	defer sqliteSchemaMu.Unlock()

	if sqliteDB == nil {
		var err error
		sqliteDB, err = gorm.Open(sqlite.Open("file:repository_tests?mode=memory&cache=shared"), &gorm.Config{})
		if err != nil {
			t.Fatalf("failed to open SQLite test db: %v", err)
		}
	}

	pending := make([]interface{}, 0, len(models))
	keys := make([]string, 0, len(models))
	for _, m := range models {
		key := fmt.Sprintf("%T", m)
		if sqliteMigrated[key] {
			continue
		}
		sqliteMigrated[key] = true
		pending = append(pending, m)
		keys = append(keys, key)
	}
	if len(pending) > 0 {
		if err := sqliteDB.AutoMigrate(pending...); err != nil {
			for _, key := range keys {
				delete(sqliteMigrated, key)
			}
			t.Fatalf("failed to auto-migrate SQLite: %v", err)
		}
	}

	tx := sqliteDB.Begin()
	if tx.Error != nil {
		t.Fatalf("failed to begin SQLite test transaction: %v", tx.Error)
	}
	t.Cleanup(func() { _ = tx.Rollback().Error })
	return tx
}

func openSQLiteDBBench(b *testing.B, models []interface{}) *gorm.DB {
	b.Helper()
	db, err := gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{
		Logger: logger.Default.LogMode(logger.Silent),
	})
	if err != nil {
		b.Fatalf("failed to open SQLite benchmark db: %v", err)
	}
	sqlDB, err := db.DB()
	if err != nil {
		b.Fatalf("failed to get SQLite benchmark db: %v", err)
	}
	b.Cleanup(func() { _ = sqlDB.Close() })
	if err := db.AutoMigrate(models...); err != nil {
		b.Fatalf("failed to auto-migrate SQLite: %v", err)
	}
	return db
}

// skipIfSQLite skips the test when the actual database is SQLite
// (use for PG-only features like ILIKE, NOW(), NULLS LAST, etc.).
// When PG is requested but unavailable, tests are also skipped since
// the fallback SQLite won't support PG-specific syntax.
func skipIfSQLite(t *testing.T) {
	t.Helper()
	if !isActuallyPG() {
		t.Skip("Skipping: this test requires PostgreSQL (ILIKE / NOW() / NULLS LAST etc.)")
	}
}