516 lines
14 KiB
Go
516 lines
14 KiB
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 = "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",
|
|
"notifications",
|
|
"usage_log",
|
|
"model_pricing",
|
|
"notes",
|
|
"audit_log",
|
|
"channel_cursors",
|
|
"messages",
|
|
"channel_models",
|
|
"channel_participants",
|
|
"channel_knowledge_bases",
|
|
"persona_knowledge_bases",
|
|
"project_notes",
|
|
"project_knowledge_bases",
|
|
"project_channels",
|
|
"kb_chunks",
|
|
"kb_documents",
|
|
"knowledge_bases",
|
|
"channels",
|
|
"projects",
|
|
"user_model_settings",
|
|
"model_catalog",
|
|
"persona_grants",
|
|
"persona_group_members",
|
|
"persona_groups",
|
|
"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": "", "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": "Chat Switchboard", "tagline": "Multi-Model AI Chat"}'::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()
|
|
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
|
|
}
|