Files
2025-12-24 19:53:45 +00:00

175 lines
4.1 KiB
Go

package integration
import (
"context"
"fmt"
"os"
"testing"
"time"
_ "github.com/lib/pq"
"github.com/atlasos/calypso/internal/common/config"
"github.com/atlasos/calypso/internal/common/database"
"github.com/atlasos/calypso/internal/common/logger"
"github.com/google/uuid"
)
// TestDB holds the test database connection
var TestDB *database.DB
var TestConfig *config.Config
var TestLogger *logger.Logger
// SetupTestDB initializes a test database connection
func SetupTestDB(t *testing.T) *database.DB {
if TestDB != nil {
return TestDB
}
// Get test database configuration from environment
dbHost := getEnv("TEST_DB_HOST", "localhost")
dbPort := getEnvInt("TEST_DB_PORT", 5432)
dbUser := getEnv("TEST_DB_USER", "calypso")
dbPassword := getEnv("TEST_DB_PASSWORD", "calypso123")
dbName := getEnv("TEST_DB_NAME", "calypso_test")
cfg := &config.Config{
Database: config.DatabaseConfig{
Host: dbHost,
Port: dbPort,
User: dbUser,
Password: dbPassword,
Database: dbName,
SSLMode: "disable",
MaxConnections: 10,
MaxIdleConns: 5,
ConnMaxLifetime: 5 * time.Minute,
},
}
db, err := database.NewConnection(cfg.Database)
if err != nil {
t.Fatalf("Failed to connect to test database: %v", err)
}
// Run migrations
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
if err := database.RunMigrations(ctx, db); err != nil {
t.Fatalf("Failed to run migrations: %v", err)
}
TestDB = db
TestConfig = cfg
if TestLogger == nil {
TestLogger = logger.NewLogger("test")
}
return db
}
// CleanupTestDB cleans up test data
func CleanupTestDB(t *testing.T) {
if TestDB == nil {
return
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
// Clean up test data (but keep schema)
tables := []string{
"sessions",
"audit_log",
"tasks",
"alerts",
"user_roles",
"role_permissions",
"users",
"scst_initiators",
"scst_luns",
"scst_targets",
"disk_repositories",
"physical_disks",
"virtual_tapes",
"virtual_tape_drives",
"virtual_tape_libraries",
"physical_tape_slots",
"physical_tape_drives",
"physical_tape_libraries",
}
for _, table := range tables {
query := fmt.Sprintf("TRUNCATE TABLE %s CASCADE", table)
if _, err := TestDB.ExecContext(ctx, query); err != nil {
// Ignore errors for tables that don't exist
t.Logf("Warning: Failed to truncate %s: %v", table, err)
}
}
}
// CreateTestUser creates a test user in the database
func CreateTestUser(t *testing.T, username, email, passwordHash string, isAdmin bool) string {
userID := uuid.New().String()
query := `
INSERT INTO users (id, username, email, password_hash, full_name, is_active)
VALUES ($1, $2, $3, $4, $5, $6)
RETURNING id
`
var id string
err := TestDB.QueryRow(query, userID, username, email, passwordHash, "Test User", true).Scan(&id)
if err != nil {
t.Fatalf("Failed to create test user: %v", err)
}
// Assign admin role if requested
if isAdmin {
roleQuery := `
INSERT INTO user_roles (user_id, role_id)
SELECT $1, id FROM roles WHERE name = 'admin'
`
if _, err := TestDB.Exec(roleQuery, id); err != nil {
t.Fatalf("Failed to assign admin role: %v", err)
}
} else {
// Assign operator role for non-admin users (gives monitoring:read permission)
roleQuery := `
INSERT INTO user_roles (user_id, role_id)
SELECT $1, id FROM roles WHERE name = 'operator'
`
if _, err := TestDB.Exec(roleQuery, id); err != nil {
// Operator role might not exist, try readonly
roleQuery = `
INSERT INTO user_roles (user_id, role_id)
SELECT $1, id FROM roles WHERE name = 'readonly'
`
if _, err := TestDB.Exec(roleQuery, id); err != nil {
t.Logf("Warning: Failed to assign role: %v", err)
}
}
}
return id
}
// Helper functions
func getEnv(key, defaultValue string) string {
if v := os.Getenv(key); v != "" {
return v
}
return defaultValue
}
func getEnvInt(key string, defaultValue int) int {
if v := os.Getenv(key); v != "" {
var result int
if _, err := fmt.Sscanf(v, "%d", &result); err == nil {
return result
}
}
return defaultValue
}