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 = "chat_switchboard_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{ "resource_grants", "group_members", "groups", "usage_log", "model_pricing", "notes", "audit_log", "channel_cursors", "messages", "channel_models", "channel_members", "channel_knowledge_bases", "persona_knowledge_bases", "kb_chunks", "kb_documents", "knowledge_bases", "channels", "user_model_settings", "model_catalog", "persona_grants", "personas", "provider_configs", "team_members", "teams", "refresh_tokens", "users", "platform_policies", "global_config", "global_settings", "folders", "attachments", "extension_user_settings", "extensions", } 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": "Chat Switchboard", "tagline": "Multi-Model AI Chat"}'), ('banner', '{"enabled": false, "text": "", "position": "both", "bg": "#007a33", "fg": "#ffffff"}'), ('banner_presets', '{}'), ('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": "Chat Switchboard", "tagline": "Multi-Model AI Chat"}'::jsonb), ('banner', '{"enabled": false, "text": "", "position": "both", "bg": "#007a33", "fg": "#ffffff"}'::jsonb), ('banner_presets', '{}'::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() if IsSQLite() { id := uuid.New().String() _, err := DB.Exec(` INSERT INTO users (id, username, email, password_hash, role) VALUES (?, ?, ?, '$2a$10$dummy.hash.for.testing.only.000000000000000000000', 'user') `, id, username, email) if err != nil { t.Fatalf("SeedTestUser: %v", err) } return id } var id string err := DB.QueryRow(` INSERT INTO users (username, email, password_hash, role) VALUES ($1, $2, '$2a$10$dummy.hash.for.testing.only.000000000000000000000', 'user') RETURNING id `, username, email).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 }