package database
import (
"context"
"crypto/rand"
"encoding/binary"
"errors"
"fmt"
"testing"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
"github.com/stretchr/testify/require"
)
const templateDBName = "mcp_registry_test_template"
func ensureTemplateDB(ctx context.Context, adminConn *pgx.Conn) error {
var exists bool
err := adminConn.QueryRow(ctx, "SELECT EXISTS(SELECT 1 FROM pg_database WHERE datname = $1)", templateDBName).Scan(&exists)
if err != nil {
return fmt.Errorf("failed to check template database: %w", err)
}
if exists {
return nil
}
_, err = adminConn.Exec(ctx, fmt.Sprintf("CREATE DATABASE %s", templateDBName))
if err != nil {
var pgErr *pgconn.PgError
if errors.As(err, &pgErr) && pgErr.Code == "23505" && pgErr.ConstraintName == "pg_database_datname_index" {
return nil
}
return fmt.Errorf("failed to create template database: %w", err)
}
templateURI := fmt.Sprintf("postgres://mcpregistry:mcpregistry@localhost:5432/%s?sslmode=disable", templateDBName)
templateDB, err := NewPostgreSQL(ctx, templateURI)
if err != nil {
return fmt.Errorf("failed to connect to template database: %w", err)
}
defer templateDB.Close()
return nil
}
func NewTestDB(t *testing.T) Database {
t.Helper()
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
adminURI := "postgres://mcpregistry:mcpregistry@localhost:5432/postgres?sslmode=disable"
adminConn, err := pgx.Connect(ctx, adminURI)
require.NoError(t, err, "Failed to connect to PostgreSQL. Make sure PostgreSQL is running via: docker-compose up -d postgres")
defer adminConn.Close(ctx)
err = ensureTemplateDB(ctx, adminConn)
require.NoError(t, err, "Failed to initialize template database")
var randomBytes [8]byte
_, err = rand.Read(randomBytes[:])
require.NoError(t, err, "Failed to generate random database id")
randomInt := binary.BigEndian.Uint64(randomBytes[:])
dbName := fmt.Sprintf("test_%d", randomInt)
_, err = adminConn.Exec(ctx, fmt.Sprintf("CREATE DATABASE %s TEMPLATE %s", dbName, templateDBName))
require.NoError(t, err, "Failed to create test database from template")
t.Cleanup(func() {
cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cleanupCancel()
_, _ = adminConn.Exec(cleanupCtx, fmt.Sprintf(
"SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '%s' AND pid <> pg_backend_pid()",
dbName,
))
_, _ = adminConn.Exec(cleanupCtx, fmt.Sprintf("DROP DATABASE IF EXISTS %s", dbName))
})
testURI := fmt.Sprintf("postgres://mcpregistry:mcpregistry@localhost:5432/%s?sslmode=disable", dbName)
db, err := NewPostgreSQL(ctx, testURI)
require.NoError(t, err, "Failed to connect to test database")
t.Cleanup(func() {
if err := db.Close(); err != nil {
t.Logf("Warning: failed to close test database connection: %v", err)
}
})
return db
}