foxhunt/tests/fixtures/test_database.rs

//! Test Database Utilities for Foxhunt HFT Trading System
//!
//! This module provides database setup, cleanup, and helper utilities
//! for testing database operations in isolation.

use std::sync::Arc;
use sqlx::{PgPool, Pool, Postgres, migrate::MigrateDatabase};
use tokio::sync::OnceCell;
use uuid::Uuid;

use super::test_config::TestConfig;

/// Test database manager for setup and cleanup
#[derive(Debug)]
pub struct TestDatabase {
    pub pool: PgPool,
    pub config: TestConfig,
    pub database_name: String,
    cleanup_on_drop: bool,
}

impl TestDatabase {
    /// Create a new test database with unique name
    pub async fn new() -> Result<Self, sqlx::Error> {
        Self::with_config(TestConfig::for_unit_tests()).await
    }

    /// Create a test database with specific configuration
    pub async fn with_config(config: TestConfig) -> Result<Self, sqlx::Error> {
        let database_name = format!("test_db_{}", Uuid::new_v4().simple());
        let base_url = Self::get_base_database_url(&config.database.url);
        let full_url = format!("{}/{}", base_url, database_name);

        // Create the test database
        Postgres::create_database(&full_url).await?;

        // Connect to the new database
        let pool = PgPool::connect(&full_url).await?;

        // Run migrations if enabled
        if config.database.auto_migrate {
            Self::run_migrations(&pool).await?;
        }

        Ok(Self {
            pool,
            config,
            database_name,
            cleanup_on_drop: true,
        })
    }

    /// Create a test database from existing pool (for shared testing)
    pub fn from_pool(pool: PgPool, config: TestConfig) -> Self {
        Self {
            pool,
            config,
            database_name: "shared_test_db".to_string(),
            cleanup_on_drop: false, // Don't cleanup shared pools
        }
    }

    /// Get the database pool
    pub fn pool(&self) -> &PgPool {
        &self.pool
    }

    /// Execute SQL and return affected rows
    pub async fn execute(&self, sql: &str) -> Result<u64, sqlx::Error> {
        let result = sqlx::query(sql).execute(&self.pool).await?;
        Ok(result.rows_affected())
    }

    /// Fetch one row
    pub async fn fetch_one<T>(&self, sql: &str) -> Result<T, sqlx::Error>
    where
        T: for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin,
    {
        sqlx::query_as::<_, T>(sql).fetch_one(&self.pool).await
    }

    /// Fetch all rows
    pub async fn fetch_all<T>(&self, sql: &str) -> Result<Vec<T>, sqlx::Error>
    where
        T: for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin,
    {
        sqlx::query_as::<_, T>(sql).fetch_all(&self.pool).await
    }

    /// Clean all test data from database (without dropping)
    pub async fn clean_test_data(&self) -> Result<(), sqlx::Error> {
        // Clean in dependency order (children first, parents last)
        let cleanup_queries = vec![
            "DELETE FROM stress_test_results WHERE portfolio_id LIKE 'TEST_%'",
            "DELETE FROM counterparty_exposure WHERE counterparty_id LIKE 'TEST_%'",
            "DELETE FROM risk_factor_exposure WHERE portfolio_id LIKE 'TEST_%'",
            "DELETE FROM portfolio_performance WHERE portfolio_id LIKE 'TEST_%'",
            "DELETE FROM positions WHERE portfolio_id LIKE 'TEST_%'",
            "DELETE FROM portfolios WHERE id LIKE 'TEST_%'",
            "DELETE FROM counterparty WHERE id LIKE 'TEST_%'",
            "DELETE FROM instruments WHERE symbol LIKE 'TEST_%'",
            "DELETE FROM stress_scenarios WHERE name LIKE 'Test %'",
            "DELETE FROM risk_calculation_jobs WHERE created_by = 'test_system'",
            "DELETE FROM market_data_feeds WHERE provider_name LIKE 'TEST_%'",
            "DELETE FROM liquidity_metrics WHERE symbol LIKE 'TEST_%'",
            "DELETE FROM economic_scenarios WHERE created_by = 'test_system'",
            "DELETE FROM risk_reports WHERE generated_by = 'test_system'",
            "DELETE FROM risk_report_templates WHERE created_by = 'test_system'",
            "DELETE FROM custom_risk_metric_results WHERE portfolio_id LIKE 'TEST_%'",
            "DELETE FROM custom_risk_metrics WHERE created_by = 'test_system'",
        ];

        for query in cleanup_queries {
            self.execute(query).await?;
        }

        Ok(())
    }

    /// Insert test data into database
    pub async fn insert_test_data(&self) -> Result<(), sqlx::Error> {
        // Insert test instruments
        self.insert_test_instruments().await?;
        
        // Insert test portfolios
        self.insert_test_portfolios().await?;
        
        // Insert test counterparties
        self.insert_test_counterparties().await?;
        
        Ok(())
    }

    /// Insert test instruments
    async fn insert_test_instruments(&self) -> Result<(), sqlx::Error> {
        use super::{ALL_TEST_SYMBOLS, builders::BatchBuilder};
        
        let instruments = BatchBuilder::create_diverse_instruments(ALL_TEST_SYMBOLS.len());
        
        for instrument in instruments {
            sqlx::query!(
                r#"
                INSERT INTO instruments (
                    id, symbol, name, instrument_type, asset_class, 
                    currency, is_active, created_at, updated_at, metadata
                ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
                ON CONFLICT (symbol) DO NOTHING
                "#,
                instrument.id,
                instrument.symbol,
                instrument.name,
                instrument.instrument_type as _,
                instrument.asset_class as _,
                instrument.currency,
                instrument.is_active,
                instrument.created_at,
                instrument.updated_at,
                instrument.metadata,
            )
            .execute(&self.pool)
            .await?;
        }
        
        Ok(())
    }

    /// Insert test portfolios
    async fn insert_test_portfolios(&self) -> Result<(), sqlx::Error> {
        use super::{TEST_PORTFOLIO_1, TEST_PORTFOLIO_2, builders::PortfolioBuilder};
        
        let portfolios = vec![
            PortfolioBuilder::new().with_id(TEST_PORTFOLIO_1).build(),
            PortfolioBuilder::new().with_id(TEST_PORTFOLIO_2).build(),
        ];
        
        for portfolio in portfolios {
            sqlx::query!(
                r#"
                INSERT INTO portfolios (
                    id, name, description, base_currency, portfolio_type,
                    inception_date, manager_id, is_active, created_at, updated_at, metadata
                ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
                ON CONFLICT (id) DO NOTHING
                "#,
                portfolio.id,
                portfolio.name,
                portfolio.description,
                portfolio.base_currency,
                portfolio.portfolio_type,
                portfolio.inception_date,
                portfolio.manager_id,
                portfolio.is_active,
                portfolio.created_at,
                portfolio.updated_at,
                portfolio.metadata,
            )
            .execute(&self.pool)
            .await?;
        }
        
        Ok(())
    }

    /// Insert test counterparties
    async fn insert_test_counterparties(&self) -> Result<(), sqlx::Error> {
        use super::builders::CounterpartyBuilder;
        
        let counterparties = vec![
            CounterpartyBuilder::new().with_id("TEST_COUNTERPARTY_001").bank().build(),
            CounterpartyBuilder::new().with_id("TEST_COUNTERPARTY_002").broker().build(),
            CounterpartyBuilder::new().with_id("TEST_COUNTERPARTY_003").exchange().build(),
        ];
        
        for counterparty in counterparties {
            sqlx::query!(
                r#"
                INSERT INTO counterparty (
                    id, name, counterparty_type, country, credit_rating,
                    is_active, netting_agreement, created_at, updated_at, metadata
                ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
                ON CONFLICT (id) DO NOTHING
                "#,
                counterparty.id,
                counterparty.name,
                counterparty.counterparty_type,
                counterparty.country,
                counterparty.credit_rating,
                counterparty.is_active,
                counterparty.netting_agreement,
                counterparty.created_at,
                counterparty.updated_at,
                counterparty.metadata,
            )
            .execute(&self.pool)
            .await?;
        }
        
        Ok(())
    }

    /// Check if database schema is ready
    pub async fn is_schema_ready(&self) -> bool {
        let result = sqlx::query!(
            "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'instruments')"
        )
        .fetch_one(&self.pool)
        .await;
        
        match result {
            Ok(row) => row.exists.unwrap_or(false),
            Err(_) => false,
        }
    }

    /// Get database statistics
    pub async fn get_stats(&self) -> Result<DatabaseStats, sqlx::Error> {
        let instruments_count = sqlx::query_scalar!(
            "SELECT COUNT(*) FROM instruments WHERE symbol LIKE 'TEST_%'"
        )
        .fetch_one(&self.pool)
        .await?
        .unwrap_or(0);

        let portfolios_count = sqlx::query_scalar!(
            "SELECT COUNT(*) FROM portfolios WHERE id LIKE 'TEST_%'"
        )
        .fetch_one(&self.pool)
        .await?
        .unwrap_or(0);

        let positions_count = sqlx::query_scalar!(
            "SELECT COUNT(*) FROM positions WHERE portfolio_id LIKE 'TEST_%'"
        )
        .fetch_one(&self.pool)
        .await?
        .unwrap_or(0);

        Ok(DatabaseStats {
            instruments_count,
            portfolios_count,
            positions_count,
        })
    }

    /// Run database migrations
    async fn run_migrations(pool: &PgPool) -> Result<(), sqlx::Error> {
        // Basic schema creation (simplified for testing)
        sqlx::query(
            r#"
            CREATE TABLE IF NOT EXISTS instruments (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                symbol VARCHAR(50) UNIQUE NOT NULL,
                name VARCHAR(255) NOT NULL,
                instrument_type VARCHAR(50) NOT NULL,
                asset_class VARCHAR(50) NOT NULL,
                currency CHAR(3) NOT NULL,
                is_active BOOLEAN DEFAULT true,
                created_at TIMESTAMPTZ DEFAULT NOW(),
                updated_at TIMESTAMPTZ DEFAULT NOW(),
                metadata JSONB DEFAULT '{}'
            )
            "#
        )
        .execute(pool)
        .await?;

        sqlx::query(
            r#"
            CREATE TABLE IF NOT EXISTS portfolios (
                id VARCHAR(50) PRIMARY KEY,
                name VARCHAR(255) NOT NULL,
                description TEXT,
                base_currency CHAR(3) NOT NULL,
                portfolio_type VARCHAR(100) NOT NULL,
                inception_date TIMESTAMPTZ NOT NULL,
                manager_id VARCHAR(100) NOT NULL,
                is_active BOOLEAN DEFAULT true,
                created_at TIMESTAMPTZ DEFAULT NOW(),
                updated_at TIMESTAMPTZ DEFAULT NOW(),
                metadata JSONB DEFAULT '{}'
            )
            "#
        )
        .execute(pool)
        .await?;

        sqlx::query(
            r#"
            CREATE TABLE IF NOT EXISTS positions (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                portfolio_id VARCHAR(50) REFERENCES portfolios(id),
                symbol VARCHAR(50) NOT NULL,
                quantity DECIMAL NOT NULL,
                average_price DECIMAL NOT NULL,
                market_price DECIMAL NOT NULL,
                market_value DECIMAL NOT NULL,
                unrealized_pnl DECIMAL NOT NULL,
                currency CHAR(3) NOT NULL,
                entry_date TIMESTAMPTZ NOT NULL,
                last_updated TIMESTAMPTZ DEFAULT NOW()
            )
            "#
        )
        .execute(pool)
        .await?;

        sqlx::query(
            r#"
            CREATE TABLE IF NOT EXISTS counterparty (
                id VARCHAR(50) PRIMARY KEY,
                name VARCHAR(255) NOT NULL,
                counterparty_type VARCHAR(100) NOT NULL,
                country CHAR(2) NOT NULL,
                credit_rating VARCHAR(10),
                is_active BOOLEAN DEFAULT true,
                netting_agreement BOOLEAN DEFAULT false,
                created_at TIMESTAMPTZ DEFAULT NOW(),
                updated_at TIMESTAMPTZ DEFAULT NOW(),
                metadata JSONB DEFAULT '{}'
            )
            "#
        )
        .execute(pool)
        .await?;

        Ok(())
    }

    /// Get base database URL without database name
    fn get_base_database_url(full_url: &str) -> String {
        if let Some(last_slash) = full_url.rfind('/') {
            full_url[..last_slash].to_string()
        } else {
            full_url.to_string()
        }
    }
}

impl Drop for TestDatabase {
    fn drop(&mut self) {
        if self.cleanup_on_drop && self.database_name != "shared_test_db" {
            // Note: We can't use async in Drop, so we schedule cleanup
            // In practice, you might want to use a cleanup service or manual cleanup
            eprintln!("TestDatabase: Consider cleaning up database: {}", self.database_name);
        }
    }
}

/// Database statistics for monitoring test data
#[derive(Debug, Clone)]
pub struct DatabaseStats {
    pub instruments_count: i64,
    pub portfolios_count: i64,
    pub positions_count: i64,
}

/// Shared test database pool for tests that don't need isolation
static SHARED_TEST_POOL: OnceCell<Arc<TestDatabase>> = OnceCell::const_new();

/// Get or create shared test database pool
pub async fn get_shared_test_db() -> Result<Arc<TestDatabase>, sqlx::Error> {
    SHARED_TEST_POOL
        .get_or_try_init(|| async {
            let config = TestConfig::for_integration_tests();
            let test_db = TestDatabase::with_config(config).await?;
            test_db.insert_test_data().await?;
            Ok(Arc::new(test_db))
        })
        .await
        .map(Arc::clone)
}

/// Test database transaction for isolated test operations
pub struct TestTransaction<'a> {
    tx: sqlx::Transaction<'a, Postgres>,
}

impl<'a> TestTransaction<'a> {
    /// Begin a new test transaction
    pub async fn begin(pool: &'a PgPool) -> Result<Self, sqlx::Error> {
        let tx = pool.begin().await?;
        Ok(Self { tx })
    }

    /// Execute SQL within transaction
    pub async fn execute(&mut self, sql: &str) -> Result<u64, sqlx::Error> {
        let result = sqlx::query(sql).execute(&mut *self.tx).await?;
        Ok(result.rows_affected())
    }

    /// Commit the transaction
    pub async fn commit(self) -> Result<(), sqlx::Error> {
        self.tx.commit().await
    }

    /// Rollback the transaction
    pub async fn rollback(self) -> Result<(), sqlx::Error> {
        self.tx.rollback().await
    }
}

/// Utility macro for test database setup
#[macro_export]
macro_rules! setup_test_db {
    () => {{
        let test_db = $crate::fixtures::test_database::TestDatabase::new().await?;
        test_db.insert_test_data().await?;
        test_db
    }};
    
    ($config:expr) => {{
        let test_db = $crate::fixtures::test_database::TestDatabase::with_config($config).await?;
        test_db.insert_test_data().await?;
        test_db
    }};
}

/// Utility macro for test transaction
#[macro_export]
macro_rules! test_transaction {
    ($db:expr, $code:block) => {{
        let mut tx = $crate::fixtures::test_database::TestTransaction::begin($db.pool()).await?;
        let result = async move { $code }.await;
        tx.rollback().await?; // Always rollback test transactions
        result
    }};
}

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

    #[tokio::test]
    async fn test_database_creation() {
        let test_db = TestDatabase::new().await.unwrap();
        assert!(test_db.is_schema_ready().await);
        assert!(test_db.database_name.starts_with("test_db_"));
    }

    #[tokio::test]
    async fn test_insert_and_clean_test_data() {
        let test_db = TestDatabase::new().await.unwrap();
        
        // Insert test data
        test_db.insert_test_data().await.unwrap();
        
        // Check data was inserted
        let stats = test_db.get_stats().await.unwrap();
        assert!(stats.instruments_count > 0);
        assert!(stats.portfolios_count > 0);
        
        // Clean test data
        test_db.clean_test_data().await.unwrap();
        
        // Check data was cleaned
        let stats_after = test_db.get_stats().await.unwrap();
        assert_eq!(stats_after.instruments_count, 0);
        assert_eq!(stats_after.portfolios_count, 0);
    }

    #[tokio::test]
    async fn test_database_stats() {
        let test_db = TestDatabase::new().await.unwrap();
        test_db.insert_test_data().await.unwrap();
        
        let stats = test_db.get_stats().await.unwrap();
        assert!(stats.instruments_count >= 0);
        assert!(stats.portfolios_count >= 0);
        assert!(stats.positions_count >= 0);
    }

    #[tokio::test]
    async fn test_transaction_rollback() {
        let test_db = TestDatabase::new().await.unwrap();
        
        // Count before transaction
        let initial_count = test_db.execute("SELECT COUNT(*) FROM instruments").await.unwrap();
        
        // Start transaction and insert data
        let mut tx = TestTransaction::begin(test_db.pool()).await.unwrap();
        tx.execute("INSERT INTO instruments (symbol, name, instrument_type, asset_class, currency) VALUES ('TX_TEST', 'Transaction Test', 'equity', 'equities', 'USD')").await.unwrap();
        tx.rollback().await.unwrap();
        
        // Count after rollback should be the same
        let final_count = test_db.execute("SELECT COUNT(*) FROM instruments").await.unwrap();
        assert_eq!(initial_count, final_count);
    }
}