base/crates/erp-server/tests/integration/test_db.rs

use sea_orm::{ConnectionTrait, Database, DatabaseBackend, Statement};
use std::sync::Arc;

use erp_server_migration::MigratorTrait;

/// 全局信号量：限制同时创建数据库的测试数量，避免 PostgreSQL 连接耗尽
static DB_SEMAPHORE: std::sync::OnceLock<Arc<tokio::sync::Semaphore>> = std::sync::OnceLock::new();

fn db_semaphore() -> &'static Arc<tokio::sync::Semaphore> {
    DB_SEMAPHORE.get_or_init(|| Arc::new(tokio::sync::Semaphore::new(4)))
}

/// 测试数据库 — 使用本地 PostgreSQL 创建隔离测试库
///
/// 连接本地 PostgreSQL（wiki/infrastructure.md 配置），为每个测试创建独立的测试数据库。
/// 不依赖 Docker/Testcontainers，与开发环境一致。
pub struct TestDb {
    db: Option<sea_orm::DatabaseConnection>,
    db_name: String,
    _permit: Option<tokio::sync::OwnedSemaphorePermit>,
}

impl TestDb {
    pub async fn new() -> Self {
        let permit = db_semaphore()
            .clone()
            .acquire_owned()
            .await
            .expect("信号量获取失败");

        let db_name = format!("erp_test_{}", uuid::Uuid::now_v7().simple());

        let admin_url = std::env::var("TEST_DB_URL")
            .unwrap_or_else(|_| "postgres://postgres:123123@localhost:5432/postgres".to_string());

        let admin_db = Database::connect(&admin_url)
            .await
            .expect("连接本地 PostgreSQL 失败，请确认服务正在运行");

        admin_db
            .execute(Statement::from_string(
                DatabaseBackend::Postgres,
                format!("CREATE DATABASE \"{}\"", db_name),
            ))
            .await
            .expect("创建测试数据库失败");

        drop(admin_db);

        // 从 admin_url 推导测试库 URL（替换路径部分）
        let test_url = if let Some(pos) = admin_url.rfind('/') {
            format!("{}/{}", &admin_url[..pos], db_name)
        } else {
            format!("postgres://postgres:123123@localhost:5432/{}", db_name)
        };
        let db = Database::connect(&test_url)
            .await
            .expect("连接测试数据库失败");

        // 运行所有迁移
        erp_server_migration::Migrator::up(&db, None)
            .await
            .expect("执行数据库迁移失败");

        Self {
            db: Some(db),
            db_name,
            _permit: Some(permit),
        }
    }

    /// 获取数据库连接引用
    pub fn db(&self) -> &sea_orm::DatabaseConnection {
        self.db.as_ref().expect("数据库连接已被释放")
    }
}

impl Drop for TestDb {
    fn drop(&mut self) {
        let db_name = self.db_name.clone();
        self.db.take();

        // 尝试在独立线程中清理，避免在 tokio runtime 内创建新 runtime
        let _ = std::thread::spawn(move || {
            let rt = tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build();
            if let Ok(rt) = rt {
                rt.block_on(async {
                    let admin_url = std::env::var("TEST_DB_URL")
                        .unwrap_or_else(|_| "postgres://postgres:123123@localhost:5432/postgres".to_string());
                    if let Ok(admin_db) = Database::connect(&admin_url).await {
                        let disconnect_sql = format!(
                            "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '{}'",
                            db_name
                        );
                        admin_db
                            .execute(Statement::from_string(DatabaseBackend::Postgres, disconnect_sql))
                            .await
                            .ok();

                        admin_db
                            .execute(Statement::from_string(
                                DatabaseBackend::Postgres,
                                format!("DROP DATABASE IF EXISTS \"{}\"", db_name),
                            ))
                            .await
                            .ok();
                    }
                });
            }
        });
    }
}