core/server/database/testhelper.go

package database

import (
	"database/sql"
	"fmt"
	"log"
	"os"
	"strings"
	"testing"

	"github.com/google/uuid"
	_ "github.com/lib/pq"
	_ "modernc.org/sqlite"
)

// TestDB holds a connection to the test database.
// Use SetupTestDB in TestMain and pass this to subtests.
var TestDB *sql.DB

const testDBName = "switchboard_core_ci"

// SetupTestDB connects to the appropriate database backend (Postgres or
// SQLite based on DB_DRIVER env), runs migrations, and sets database.DB
// so handlers can use it.
//
// Call in TestMain:
//
//	func TestMain(m *testing.M) {
//	    teardown := database.SetupTestDB()
//	    code := m.Run()
//	    teardown()
//	    os.Exit(code)
//	}
//
// For Postgres: requires PGHOST+PGUSER or TEST_DATABASE_URL.
// For SQLite:   set DB_DRIVER=sqlite (uses temp file, no external deps).
//
// Returns a cleanup function.
func SetupTestDB() func() {
	driver := os.Getenv("DB_DRIVER")
	if driver == "sqlite" {
		return setupSQLiteTestDB()
	}
	return setupPostgresTestDB()
}

// ── SQLite test setup ───────────────────────

func setupSQLiteTestDB() func() {
	CurrentDialect = DialectSQLite

	// Use a temp file so multiple connections share the same DB.
	tmpFile, err := os.CreateTemp("", "switchboard-test-*.db")
	if err != nil {
		log.Printf("⚠ Cannot create temp SQLite DB: %v — skipping DB tests", err)
		return func() {}
	}
	dbPath := tmpFile.Name()
	tmpFile.Close()

	DB, err = sql.Open("sqlite", dbPath+"?_pragma=journal_mode(WAL)&_pragma=foreign_keys(1)&_pragma=busy_timeout(5000)")
	if err != nil {
		log.Printf("⚠ Cannot open SQLite DB: %v — skipping DB tests", err)
		os.Remove(dbPath)
		return func() {}
	}
	if err := DB.Ping(); err != nil {
		DB.Close()
		DB = nil
		log.Printf("⚠ Cannot ping SQLite DB: %v — skipping DB tests", err)
		os.Remove(dbPath)
		return func() {}
	}

	// SQLite: single writer
	DB.SetMaxOpenConns(1)

	TestDB = DB

	// Run SQLite migrations
	if err := Migrate(); err != nil {
		DB.Close()
		DB = nil
		TestDB = nil
		os.Remove(dbPath)
		log.Fatalf("❌ SQLite migration failed on test DB: %v", err)
	}

	log.Printf("✓ SQLite test database ready: %s", dbPath)

	return func() {
		if DB != nil {
			DB.Close()
			DB = nil
			TestDB = nil
		}
		os.Remove(dbPath)
		log.Printf("✓ Removed SQLite test database: %s", dbPath)
	}
}

// ── Postgres test setup (original) ──────────

func setupPostgresTestDB() func() {
	CurrentDialect = DialectPostgres

	mainDSN := os.Getenv("TEST_DATABASE_URL")
	host := envOr("PGHOST", "")
	port := envOr("PGPORT", "5432")
	user := envOr("PGUSER", "")
	pass := envOr("PGPASSWORD", "")

	if mainDSN == "" {
		if host == "" || user == "" {
			log.Println("⚠ TEST_DATABASE_URL / PGHOST+PGUSER not set — skipping DB tests")
			return func() {}
		}
		mainDSN = fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=postgres sslmode=disable",
			host, port, user, pass)
	}

	adminDB, err := sql.Open("postgres", mainDSN)
	if err != nil {
		log.Printf("⚠ Cannot connect to admin DB: %v — skipping DB tests", err)
		return func() {}
	}
	if err := adminDB.Ping(); err != nil {
		adminDB.Close()
		log.Printf("⚠ Cannot ping admin DB: %v — skipping DB tests", err)
		return func() {}
	}

	createdByUs := false
	var dbExists bool
	err = adminDB.QueryRow("SELECT 1 FROM pg_database WHERE datname = $1", testDBName).Scan(&dbExists)
	if err != nil {
		if _, err := adminDB.Exec(fmt.Sprintf("CREATE DATABASE %s", testDBName)); err != nil {
			log.Printf("⚠ Cannot CREATE DATABASE %s (will try to connect to existing): %v", testDBName, err)
		} else {
			createdByUs = true
			log.Printf("✓ Created test database: %s", testDBName)
		}
	} else {
		log.Printf("✓ Test database %s already exists (keeping extensions)", testDBName)
	}

	var testDSN string
	if host != "" {
		testDSN = fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=disable",
			host, port, user, pass, testDBName)
	} else {
		testDSN = replaceDBName(mainDSN, testDBName)
	}

	DB, err = sql.Open("postgres", testDSN)
	if err != nil {
		adminDB.Close()
		log.Printf("⚠ Cannot connect to test DB: %v — skipping DB tests", err)
		return func() {}
	}
	if err := DB.Ping(); err != nil {
		DB.Close()
		DB = nil
		adminDB.Close()
		log.Printf("⚠ Cannot ping test DB: %v — skipping DB tests", err)
		return func() {}
	}
	TestDB = DB

	for _, ext := range []string{"uuid-ossp", "pgcrypto", "vector"} {
		DB.Exec(fmt.Sprintf(`CREATE EXTENSION IF NOT EXISTS "%s"`, ext))
	}

	// Wipe all tables so migrations apply cleanly
	DB.Exec(`
		DO $$ DECLARE r RECORD;
		BEGIN
			FOR r IN (SELECT tablename FROM pg_tables WHERE schemaname = 'public') LOOP
				EXECUTE 'DROP TABLE IF EXISTS public.' || quote_ident(r.tablename) || ' CASCADE';
			END LOOP;
		END $$;
	`)

	if err := Migrate(); err != nil {
		DB.Close()
		DB = nil
		TestDB = nil
		if createdByUs {
			adminDB.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %s", testDBName))
		}
		adminDB.Close()
		log.Fatalf("❌ Migration failed on test DB: %v", err)
	}

	return func() {
		if DB != nil {
			DB.Close()
			DB = nil
			TestDB = nil
		}
		if createdByUs {
			adminDB.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %s", testDBName))
			log.Printf("✓ Dropped test database: %s", testDBName)
		}
		adminDB.Close()
	}
}

// ── Dialect helpers ─────────────────────────

// PH returns a placeholder for the Nth parameter (1-indexed).
// Postgres: $1, $2, ...   SQLite: ?, ?, ...
// Exported for use by test files in other packages (e.g. handlers).
func PH(n int) string {
	if IsSQLite() {
		return "?"
	}
	return fmt.Sprintf("$%d", n)
}

// placeholders returns "?, ?, ?" or "$1, $2, $3" for n params.
func placeholders(n int) string {
	parts := make([]string, n)
	for i := range parts {
		parts[i] = PH(i + 1)
	}
	return strings.Join(parts, ", ")
}

// nowSQL returns the SQL expression for "current timestamp".
func nowSQL() string {
	if IsSQLite() {
		return "datetime('now')"
	}
	return "NOW()"
}

// ── Shared test helpers ─────────────────────

// RequireTestDB skips a test if no test database is available.
func RequireTestDB(t *testing.T) {
	t.Helper()
	if DB == nil || TestDB == nil {
		t.Skip("requires TEST_DATABASE_URL or PGHOST+PGUSER environment variables (or DB_DRIVER=sqlite)")
	}
}

// TruncateAll truncates all application tables for test isolation.
func TruncateAll(t *testing.T) {
	t.Helper()
	if DB == nil {
		return
	}

	tables := []string{
		// HA
		"rate_limit_counters",
		"ws_tickets",
		// Workflows
		"workflow_versions",
		"workflow_stages",
		"workflows",
		// Extensions
		"ext_data_tables",
		"extension_permissions",
		"ext_dependencies",
		"ext_connections",
		"package_user_settings",
		"packages",
		// Resource grants & groups
		"resource_grants",
		"group_members",
		"groups",
		// Notifications
		"notification_preferences",
		"notifications",
		// Audit
		"audit_log",
		// Presence
		"user_presence",
		// Teams & users
		"team_members",
		"teams",
		"refresh_tokens",
		"user_presence",
		"users",
		// Config & packages
		"platform_policies",
		"global_settings",
		"package_user_settings",
		"packages",
		"oidc_auth_state",
	}

	if IsSQLite() {
		DB.Exec("PRAGMA foreign_keys = OFF")
		for _, table := range tables {
			DB.Exec(fmt.Sprintf("DELETE FROM %s", table))
		}
		DB.Exec("PRAGMA foreign_keys = ON")
	} else {
		for _, table := range tables {
			DB.Exec(fmt.Sprintf("TRUNCATE TABLE %s CASCADE", table))
		}
	}

	// Re-seed default config rows.
	if IsSQLite() {
		DB.Exec(`
			INSERT INTO global_settings (key, value) VALUES
				('registration', '{"enabled": true}'),
				('site', '{"name": "Switchboard Core", "tagline": "Self-hosted extension platform"}'),
				('banner', '{"enabled": false, "text": "", "bg": "#007a33", "fg": "#ffffff"}'),
				('model_roles', '{"utility":{"primary":null,"fallback":null},"embedding":{"primary":null,"fallback":null},"generation":{"primary":null,"fallback":null}}')
			ON CONFLICT (key) DO NOTHING
		`)
		DB.Exec(`
			INSERT INTO platform_policies (key, value) VALUES
				('allow_user_byok',         'false'),
				('allow_user_personas',     'false'),
				('allow_raw_model_access',  'true'),
				('allow_registration',      'true'),
				('default_user_active',     'false'),
				('allow_team_providers',    'true')
			ON CONFLICT (key) DO NOTHING
		`)
	} else {
		DB.Exec(`
			INSERT INTO global_settings (key, value) VALUES
				('registration', '{"enabled": true}'::jsonb),
				('site', '{"name": "Switchboard Core", "tagline": "Self-hosted extension platform"}'::jsonb),
				('banner', '{"enabled": false, "text": "", "bg": "#007a33", "fg": "#ffffff"}'::jsonb),
				('model_roles', '{
					"utility":    { "primary": null, "fallback": null },
					"embedding":  { "primary": null, "fallback": null },
					"generation": { "primary": null, "fallback": null }
				}'::jsonb)
			ON CONFLICT (key) DO NOTHING
		`)
		DB.Exec(`
			INSERT INTO platform_policies (key, value) VALUES
				('allow_user_byok',         'false'),
				('allow_user_personas',     'false'),
				('allow_raw_model_access',  'true'),
				('allow_registration',      'true'),
				('default_user_active',     'false'),
				('allow_team_providers',    'true')
			ON CONFLICT (key) DO NOTHING
		`)
	}
}

// SeedTestUser creates a test user and returns the user ID.
func SeedTestUser(t *testing.T, username, email string) string {
	t.Helper()
	handle := strings.ToLower(strings.ReplaceAll(username, " ", "-"))
	if IsSQLite() {
		id := uuid.New().String()
		_, err := DB.Exec(`
			INSERT INTO users (id, username, email, password_hash, role, auth_source, handle)
			VALUES (?, ?, ?, '$2a$10$dummy.hash.for.testing.only.000000000000000000000', 'user', 'builtin', ?)
		`, id, username, email, handle)
		if err != nil {
			t.Fatalf("SeedTestUser: %v", err)
		}
		return id
	}
	var id string
	err := DB.QueryRow(`
		INSERT INTO users (username, email, password_hash, role, auth_source, handle)
		VALUES ($1, $2, '$2a$10$dummy.hash.for.testing.only.000000000000000000000', 'user', 'builtin', $3)
		RETURNING id
	`, username, email, handle).Scan(&id)
	if err != nil {
		t.Fatalf("SeedTestUser: %v", err)
	}
	return id
}

// SeedTestChannel creates a test channel owned by userID and returns the channel ID.
func SeedTestChannel(t *testing.T, userID, title string) string {
	t.Helper()
	if IsSQLite() {
		id := uuid.New().String()
		_, err := DB.Exec(`INSERT INTO channels (id, user_id, title) VALUES (?, ?, ?)`, id, userID, title)
		if err != nil {
			t.Fatalf("SeedTestChannel: %v", err)
		}
		return id
	}
	var id string
	err := DB.QueryRow(`INSERT INTO channels (user_id, title) VALUES ($1, $2) RETURNING id`, userID, title).Scan(&id)
	if err != nil {
		t.Fatalf("SeedTestChannel: %v", err)
	}
	return id
}

// SeedTestTeam creates a test team and returns the team ID.
func SeedTestTeam(t *testing.T, name, createdBy string) string {
	t.Helper()
	if IsSQLite() {
		id := uuid.New().String()
		_, err := DB.Exec(`INSERT INTO teams (id, name, created_by) VALUES (?, ?, ?)`, id, name, createdBy)
		if err != nil {
			t.Fatalf("SeedTestTeam: %v", err)
		}
		return id
	}
	var id string
	err := DB.QueryRow(`INSERT INTO teams (name, created_by) VALUES ($1, $2) RETURNING id`, name, createdBy).Scan(&id)
	if err != nil {
		t.Fatalf("SeedTestTeam: %v", err)
	}
	return id
}

// SeedTestTeamMember adds a user to a team.
func SeedTestTeamMember(t *testing.T, teamID, userID, role string) {
	t.Helper()
	if IsSQLite() {
		_, err := DB.Exec(`INSERT INTO team_members (id, team_id, user_id, role) VALUES (?, ?, ?, ?)`,
			uuid.New().String(), teamID, userID, role)
		if err != nil {
			t.Fatalf("SeedTestTeamMember: %v", err)
		}
		return
	}
	q := fmt.Sprintf(`
		INSERT INTO team_members (team_id, user_id, role)
		VALUES (%s, %s, %s)
	`, PH(1), PH(2), PH(3))
	_, err := DB.Exec(q, teamID, userID, role)
	if err != nil {
		t.Fatalf("SeedTestTeamMember: %v", err)
	}
}

// SeedTestGroup creates a test group and returns the group ID.
func SeedTestGroup(t *testing.T, name, scope string, teamID *string, createdBy string) string {
	t.Helper()
	var teamArg interface{}
	if teamID != nil {
		teamArg = *teamID
	}
	if IsSQLite() {
		id := uuid.New().String()
		_, err := DB.Exec(`INSERT INTO groups (id, name, scope, team_id, created_by) VALUES (?, ?, ?, ?, ?)`,
			id, name, scope, teamArg, createdBy)
		if err != nil {
			t.Fatalf("SeedTestGroup: %v", err)
		}
		return id
	}
	var id string
	err := DB.QueryRow(`INSERT INTO groups (name, scope, team_id, created_by) VALUES ($1, $2, $3, $4) RETURNING id`,
		name, scope, teamArg, createdBy).Scan(&id)
	if err != nil {
		t.Fatalf("SeedTestGroup: %v", err)
	}
	return id
}

// SeedGroupMember adds a user to a group.
func SeedGroupMember(t *testing.T, groupID, userID, addedBy string) {
	t.Helper()
	if IsSQLite() {
		_, err := DB.Exec(`INSERT INTO group_members (id, group_id, user_id, added_by) VALUES (?, ?, ?, ?)`,
			uuid.New().String(), groupID, userID, addedBy)
		if err != nil {
			t.Fatalf("SeedGroupMember: %v", err)
		}
		return
	}
	q := fmt.Sprintf(`
		INSERT INTO group_members (group_id, user_id, added_by)
		VALUES (%s, %s, %s)
	`, PH(1), PH(2), PH(3))
	_, err := DB.Exec(q, groupID, userID, addedBy)
	if err != nil {
		t.Fatalf("SeedGroupMember: %v", err)
	}
}

// ── Internal helpers ────────────────────────

func envOr(key, fallback string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return fallback
}

func replaceDBName(dsn, newDB string) string {
	if strings.Contains(dsn, "dbname=") {
		parts := strings.Fields(dsn)
		for i, p := range parts {
			if strings.HasPrefix(p, "dbname=") {
				parts[i] = "dbname=" + newDB
				return strings.Join(parts, " ")
			}
		}
	}
	if strings.Contains(dsn, "://") {
		idx := strings.LastIndex(dsn, "/")
		qIdx := strings.Index(dsn, "?")
		if qIdx > idx {
			return dsn[:idx+1] + newDB + dsn[qIdx:]
		}
		return dsn[:idx+1] + newDB
	}
	return dsn + " dbname=" + newDB
}