foxhunt/testing/test-common/src/fixtures/database.rs

//! Database test fixtures
//!
//! Provides shared database setup/teardown utilities used across 39+ test files.

use anyhow::Result;
use once_cell::sync::OnceCell;
use sqlx::{postgres::PgPoolOptions, PgPool};
use std::sync::Arc;
use tokio::sync::Mutex;

static DB_POOL: OnceCell<Arc<Mutex<PgPool>>> = OnceCell::new();

/// Test database wrapper with automatic setup and cleanup
pub struct TestDb {
    pool: PgPool,
    schema_name: String,
}

impl TestDb {
    /// Create a new test database instance
    ///
    /// Uses the DATABASE_URL environment variable or falls back to default.
    /// Creates an isolated schema for test isolation.
    pub async fn new() -> Self {
        let pool = get_test_pool().await;
        let schema_name = format!("test_{}", uuid::Uuid::new_v4().to_string().replace('-', "_"));

        // Create isolated schema
        sqlx::query(&format!("CREATE SCHEMA {}", schema_name))
            .execute(&pool)
            .await
            .expect("Failed to create test schema");

        Self {
            pool,
            schema_name,
        }
    }

    /// Create a new test database with migrations applied
    pub async fn with_migrations() -> Self {
        let db = Self::new().await;

        // Set search path to test schema
        sqlx::query(&format!("SET search_path TO {}", db.schema_name))
            .execute(&db.pool)
            .await
            .expect("Failed to set search path");

        // Run migrations
        sqlx::migrate!("../../migrations")
            .run(&db.pool)
            .await
            .expect("Failed to run migrations");

        db
    }

    /// Get a reference to the connection pool
    pub fn pool(&self) -> &PgPool {
        &self.pool
    }

    /// Execute a query in the test schema
    pub async fn execute(&self, sql: &str) -> Result<sqlx::postgres::PgQueryResult> {
        sqlx::query(&format!("SET search_path TO {}; {}", self.schema_name, sql))
            .execute(&self.pool)
            .await
            .map_err(Into::into)
    }

    /// Begin a transaction for test isolation
    pub async fn begin(&self) -> Result<sqlx::Transaction<'_, sqlx::Postgres>> {
        self.pool.begin().await.map_err(Into::into)
    }
}

impl Drop for TestDb {
    fn drop(&mut self) {
        // Clean up test schema
        let pool = self.pool.clone();
        let schema_name = self.schema_name.clone();

        tokio::spawn(async move {
            let _ = sqlx::query(&format!("DROP SCHEMA IF EXISTS {} CASCADE", schema_name))
                .execute(&pool)
                .await;
        });
    }
}

/// Get or create the shared test database pool
///
/// Pattern found in 39 test files:
/// - services/ml_training_service/tests/checkpoint_manager_tests.rs
/// - services/ml_training_service/tests/job_tracker_test.rs
/// - database/src/test_helpers.rs
/// - And 36+ more
async fn get_test_pool() -> PgPool {
    let pool = DB_POOL.get_or_init(|| {
        let database_url = std::env::var("DATABASE_URL").unwrap_or_else(|_| {
            "postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt".to_string()
        });

        Arc::new(Mutex::new(
            tokio::task::block_in_place(|| {
                tokio::runtime::Handle::current().block_on(async {
                    PgPoolOptions::new()
                        .max_connections(5)
                        .connect(&database_url)
                        .await
                        .expect("Failed to create test pool")
                })
            })
        ))
    });

    pool.lock().await.clone()
}

/// Setup test data in the database
///
/// Pattern found in data_loader_integration.rs
pub async fn setup_test_data(pool: &PgPool, symbol_prefix: &str) -> Result<()> {
    // Clean existing test data
    sqlx::query(&format!(
        "DELETE FROM market_events WHERE symbol LIKE '{}%'",
        symbol_prefix
    ))
    .execute(pool)
    .await?;

    sqlx::query(&format!(
        "DELETE FROM trade_executions WHERE symbol LIKE '{}%'",
        symbol_prefix
    ))
    .execute(pool)
    .await?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_db_creation() {
        let db = TestDb::new().await;
        assert!(db.pool().is_closed() == false);
    }

    #[tokio::test]
    async fn test_db_with_migrations() {
        let db = TestDb::with_migrations().await;
        let result = db.execute("SELECT 1").await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_schema_isolation() {
        let db1 = TestDb::new().await;
        let db2 = TestDb::new().await;

        // Different schemas should be isolated
        assert_ne!(db1.schema_name, db2.schema_name);
    }
}