Agent-System/internal/db/testdb.go

package db

import (
	"context"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"fmt"
	"net/url"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"testing"

	"github.com/jackc/pgx/v5/pgxpool"
)

var (
	migrateOnceMu sync.Mutex
	migratedDSNs  = map[string]error{}
)

// RunMigrationsOnceForTest 同一 DSN 只迁一次，避免并行包抢同一隔离库。
func RunMigrationsOnceForTest(ctx context.Context, pool *pgxpool.Pool, dsn string) error {
	key := strings.TrimSpace(dsn)
	migrateOnceMu.Lock()
	defer migrateOnceMu.Unlock()
	if err, ok := migratedDSNs[key]; ok {
		return err
	}
	err := RunMigrations(ctx, pool, ResolveMigrationsDir())
	migratedDSNs[key] = err
	return err
}

const testDatabaseEnv = "TEST_DATABASE_URL"

// RequireTestDatabaseURL 集成测试必须显式提供隔离测试库，禁止默认连开发库。
// 返回按调用方包派生的独立数据库（<base>_<pkgHash>）：包内共享、包间隔离，
// 消除 go test 并行包对同一库的迁移/种子互相干扰（2026-09-29 issue #2）。
func RequireTestDatabaseURL(t *testing.T) string {
	t.Helper()
	dsn := strings.TrimSpace(os.Getenv(testDatabaseEnv))
	if dsn == "" {
		t.Fatalf("%s must be set for PostgreSQL integration tests (refusing default dev database)", testDatabaseEnv)
	}
	return derivedPackageDSN(t, dsn)
}

// SkipIfNoTestDatabaseURL 无隔离测试库时跳过（单元测试包内可选 PG 用例）。
// 同样返回按包派生库。
func SkipIfNoTestDatabaseURL(t *testing.T) string {
	t.Helper()
	dsn := strings.TrimSpace(os.Getenv(testDatabaseEnv))
	if dsn == "" {
		t.Skipf("%s not set; skipping PostgreSQL integration test", testDatabaseEnv)
	}
	return derivedPackageDSN(t, dsn)
}

var (
	pkgDBMu    sync.Mutex
	pkgDBReady = map[string]string{} // 派生 DSN 缓存（进程内）
)

func derivedPackageDSN(t *testing.T, baseDSN string) string {
	t.Helper()
	_, file, _, ok := runtime.Caller(2) // 跳过 helper 自身，取调用测试文件
	if !ok {
		return baseDSN
	}
	pkgDir := filepath.ToSlash(filepath.Dir(file))
	// 仅取模块内相对路径段，避免绝对路径差异影响哈希
	if i := strings.Index(pkgDir, "internal/"); i >= 0 {
		pkgDir = pkgDir[i:]
	} else if i := strings.Index(pkgDir, "cmd/"); i >= 0 {
		pkgDir = pkgDir[i:]
	}
	sum := sha256.Sum256([]byte(pkgDir))
	suffix := hex.EncodeToString(sum[:4])

	u, err := url.Parse(baseDSN)
	if err != nil || u.Path == "" {
		return baseDSN // 非标准 DSN：退回原值，行为同旧版
	}
	baseDB := strings.TrimPrefix(u.Path, "/")
	if baseDB == "" {
		return baseDSN
	}
	sanitized := strings.Map(func(r rune) rune {
		switch {
		case r >= 'a' && r <= 'z', r >= '0' && r <= '9':
			return r
		default:
			return '_'
		}
	}, strings.ToLower(baseDB))
	name := fmt.Sprintf("%s_%s", sanitized, suffix)
	if len(name) > 63 {
		name = name[:63]
	}

	u2 := *u
	u2.Path = "/" + name

	pkgDBMu.Lock()
	cached, done := pkgDBReady[name]
	pkgDBMu.Unlock()
	if done {
		return cached
	}

	if err := ensureTestDatabase(baseDSN, name); err != nil {
		t.Fatalf("prepare per-package test database %s: %v", name, err)
	}
	dsn := u2.String()
	pkgDBMu.Lock()
	pkgDBReady[name] = dsn
	pkgDBMu.Unlock()
	return dsn
}

// ensureTestDatabase 在基库实例上幂等创建派生库（advisory lock 防并行包竞态）。
func ensureTestDatabase(baseDSN, name string) error {
	admindb, err := sql.Open("pgx", baseDSN)
	if err != nil {
		return err
	}
	defer admindb.Close()
	var exists bool
	if err := admindb.QueryRow(`SELECT EXISTS(SELECT 1 FROM pg_database WHERE datname=$1)`, name).Scan(&exists); err != nil {
		return err
	}
	if exists {
		return nil
	}
	if _, err := admindb.Exec(`SELECT pg_advisory_lock(hashtext($1))`, name); err != nil {
		return err
	}
	defer admindb.Exec(`SELECT pg_advisory_unlock(hashtext($1))`, name)
	if err := admindb.QueryRow(`SELECT EXISTS(SELECT 1 FROM pg_database WHERE datname=$1)`, name).Scan(&exists); err != nil {
		return err
	}
	if exists {
		return nil
	}
	_, err = admindb.Exec(fmt.Sprintf(`CREATE DATABASE %q`, name))
	return err
}

// RefuseSharedEnvDSN fails if the DSN targets a shared env database.
func RefuseSharedEnvDSN(t *testing.T, dsn string) {
	t.Helper()
	if strings.Contains(dsn, "5433") || strings.Contains(dsn, "agent_system_dev") ||
		(strings.Contains(dsn, "5434") && strings.Contains(dsn, "agent_system_test")) ||
		strings.Contains(dsn, "5435") {
		t.Fatal("TEST_DATABASE_URL points at a shared env port; refuse agent_system_dev/test/main")
	}
}