277 lines
7.6 KiB
Go
277 lines
7.6 KiB
Go
package database
|
|
|
|
import (
|
|
"database/sql"
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
|
|
_ "github.com/lib/pq"
|
|
)
|
|
|
|
// 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 PostgreSQL, creates the CI test database (or
|
|
// connects to an existing one created by CI bootstrap), runs all
|
|
// 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)
|
|
// }
|
|
//
|
|
// Requires env: TEST_DATABASE_URL (full DSN to maintenance DB, e.g. postgres)
|
|
// OR individual: PGHOST, PGPORT, PGUSER, PGPASSWORD
|
|
//
|
|
// Returns a cleanup function that drops the test database.
|
|
func SetupTestDB() func() {
|
|
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)
|
|
}
|
|
|
|
// Try to connect to admin/maintenance DB for create/drop operations
|
|
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() {}
|
|
}
|
|
|
|
// Attempt to create the test DB. If the user lacks CREATE DATABASE
|
|
// permissions (e.g. CI already created it via admin bootstrap step),
|
|
// that's fine — we'll just connect to the existing one.
|
|
createdByUs := false
|
|
adminDB.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %s", testDBName))
|
|
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)
|
|
}
|
|
|
|
// Build DSN for the test DB
|
|
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)
|
|
}
|
|
|
|
// Connect to test DB
|
|
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
|
|
|
|
// Ensure required extensions (may already exist from CI bootstrap)
|
|
for _, ext := range []string{"uuid-ossp", "pgcrypto"} {
|
|
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 $$;
|
|
`)
|
|
|
|
// Run all migrations
|
|
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 cleanup
|
|
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()
|
|
}
|
|
}
|
|
|
|
// 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")
|
|
}
|
|
}
|
|
|
|
// TruncateAll truncates all application tables for test isolation.
|
|
// Faster than recreating the DB for each test.
|
|
func TruncateAll(t *testing.T) {
|
|
t.Helper()
|
|
if DB == nil {
|
|
return
|
|
}
|
|
// Order matters due to foreign keys — truncate with CASCADE
|
|
tables := []string{
|
|
"usage_log",
|
|
"model_pricing",
|
|
"notes",
|
|
"audit_log",
|
|
"channel_cursors",
|
|
"messages",
|
|
"channel_models",
|
|
"channel_members",
|
|
"channels",
|
|
"user_model_settings",
|
|
"model_catalog",
|
|
"persona_grants",
|
|
"personas",
|
|
"provider_configs",
|
|
"team_members",
|
|
"teams",
|
|
"refresh_tokens",
|
|
"users",
|
|
}
|
|
for _, table := range tables {
|
|
DB.Exec(fmt.Sprintf("TRUNCATE TABLE %s CASCADE", table))
|
|
}
|
|
|
|
// Re-seed global_settings — CASCADE from users wipes rows with updated_by FK.
|
|
// Re-run the seed SQL from migrations 001 + 004.
|
|
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
|
|
`)
|
|
|
|
// Re-seed platform_policies — also wiped by CASCADE from users truncation
|
|
// (platform_policies.updated_by REFERENCES users(id)).
|
|
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()
|
|
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()
|
|
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
|
|
}
|
|
|
|
// ── Helpers ─────────────────────────────────
|
|
|
|
func envOr(key, fallback string) string {
|
|
if v := os.Getenv(key); v != "" {
|
|
return v
|
|
}
|
|
return fallback
|
|
}
|
|
|
|
// replaceDBName swaps the dbname in a DSN string.
|
|
func replaceDBName(dsn, newDB string) string {
|
|
// Handle key=value format
|
|
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, " ")
|
|
}
|
|
}
|
|
}
|
|
// Handle URL format: postgres://user:pass@host/olddb?...
|
|
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
|
|
}
|
|
// Fallback: append dbname
|
|
return dsn + " dbname=" + newDB
|
|
}
|