Changeset 0.23.0 (#153)
This commit is contained in:
513
database/testhelper.go
Normal file
513
database/testhelper.go
Normal file
@@ -0,0 +1,513 @@
|
||||
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",
|
||||
"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
|
||||
}
|
||||
Reference in New Issue
Block a user