foxhunt/services/ml_training_service/tests/data_loader_integration.rs

//! Integration tests for HistoricalDataLoader
//!
//! These tests verify the data loading pipeline with a real PostgreSQL database.
//! They require a test database instance to be running.
//!
//! ## Running Tests
//!
//! ```bash
//! # Set up test database
//! export TEST_DATABASE_URL="postgresql://postgres:password@localhost:5432/foxhunt_test"
//!
//! # Run integration tests
//! cargo test --test data_loader_integration -- --test-threads=1
//! ```
//!
//! ## Test Database Setup
//!
//! The tests use a dedicated test database to avoid conflicts with production data.
//! Before running, ensure:
//! 1. PostgreSQL is running
//! 2. Test database exists
//! 3. Migrations have been applied
//!
//! ```sql
//! CREATE DATABASE foxhunt_test;
//! ```

use chrono::Utc;
use ml_training_service::data_config::{
    CacheConfig, DataSourceType, DataValidationConfig, DatabaseConfig, DatabaseTables,
    FeatureExtractionConfig, TimeRangeConfig, TrainingDataSourceConfig,
};
use ml_training_service::data_loader::HistoricalDataLoader;
use ml_training_service::schema_types::{MarketEvent, OrderBookSnapshot, TradeExecution};
use sqlx::PgPool;
use std::env;

/// Get test database URL from environment
fn get_test_database_url() -> String {
    env::var("TEST_DATABASE_URL")
        .unwrap_or_else(|_| "postgresql://postgres:password@localhost:5432/foxhunt_test".to_string())
}

/// Create test database connection pool
async fn create_test_pool() -> Result<PgPool, sqlx::Error> {
    let database_url = get_test_database_url();
    sqlx::postgres::PgPoolOptions::new()
        .max_connections(5)
        .connect(&database_url)
        .await
}

/// Setup test database with sample data
async fn setup_test_data(pool: &PgPool) -> Result<(), sqlx::Error> {
    // Clean existing test data
    sqlx::query("DELETE FROM market_events WHERE symbol LIKE 'TEST%'")
        .execute(pool)
        .await?;
    sqlx::query("DELETE FROM trade_executions WHERE symbol LIKE 'TEST%'")
        .execute(pool)
        .await?;
    sqlx::query("DELETE FROM order_book_snapshots WHERE symbol LIKE 'TEST%'")
        .execute(pool)
        .await?;

    // Insert test order book snapshots
    for i in 0..100 {
        let timestamp = Utc::now() - chrono::Duration::minutes(100 - i);
        let price = 100.0 + (i as f64 * 0.1);

        sqlx::query(
            r#"
            INSERT INTO order_book_snapshots
            (timestamp, symbol, best_bid, best_ask, bid_volume, ask_volume, spread_bps, mid_price, imbalance)
            VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
            "#,
        )
        .bind(timestamp)
        .bind("TEST_SYMBOL")
        .bind(rust_decimal::Decimal::from_f64_retain(price - 0.01).unwrap())
        .bind(rust_decimal::Decimal::from_f64_retain(price + 0.01).unwrap())
        .bind(rust_decimal::Decimal::new(1000, 0))
        .bind(rust_decimal::Decimal::new(800, 0))
        .bind(2i32)
        .bind(rust_decimal::Decimal::from_f64_retain(price).unwrap())
        .bind(0.111)
        .execute(pool)
        .await?;
    }

    // Insert test trade executions
    for i in 0..50 {
        let timestamp = Utc::now() - chrono::Duration::minutes(50 - i);
        let price = 100.0 + (i as f64 * 0.2);

        sqlx::query(
            r#"
            INSERT INTO trade_executions
            (timestamp, symbol, price, quantity, side)
            VALUES ($1, $2, $3, $4, $5)
            "#,
        )
        .bind(timestamp)
        .bind("TEST_SYMBOL")
        .bind(rust_decimal::Decimal::from_f64_retain(price).unwrap())
        .bind(rust_decimal::Decimal::new(100, 0))
        .bind(if i % 2 == 0 { "buy" } else { "sell" })
        .execute(pool)
        .await?;
    }

    // Insert test market events
    for i in 0..10 {
        let timestamp = Utc::now() - chrono::Duration::hours(10 - i);

        sqlx::query(
            r#"
            INSERT INTO market_events
            (timestamp, event_type, symbol, title, impact_score, sentiment)
            VALUES ($1, $2, $3, $4, $5, $6)
            "#,
        )
        .bind(timestamp)
        .bind("news")
        .bind("TEST_SYMBOL")
        .bind(format!("Test Event {}", i))
        .bind(0.5)
        .bind(0.3)
        .execute(pool)
        .await?;
    }

    Ok(())
}

/// Create test training data configuration
fn create_test_config() -> TrainingDataSourceConfig {
    let database_url = get_test_database_url();

    TrainingDataSourceConfig {
        source_type: DataSourceType::Historical,
        database: Some(DatabaseConfig {
            connection_url: database_url,
            max_connections: 5,
            query_timeout_secs: 30,
            tables: DatabaseTables::default(),
        }),
        s3: None,
        time_range: TimeRangeConfig {
            start: Some(Utc::now() - chrono::Duration::hours(2)),
            end: Some(Utc::now()),
            duration_days: None,
            train_split: 0.8,
        },
        symbols: vec!["TEST_SYMBOL".to_string()],
        features: FeatureExtractionConfig::default(),
        validation: DataValidationConfig {
            min_samples: 10,
            max_missing_ratio: 0.2,
            enable_outlier_detection: true,
            outlier_threshold: 3.0,
        },
        cache: CacheConfig::default(),
    }
}

#[tokio::test]
#[ignore] // Requires test database setup
async fn test_load_historical_data() {
    // Setup
    let pool = create_test_pool().await.expect("Failed to create test pool");
    setup_test_data(&pool).await.expect("Failed to setup test data");

    let config = create_test_config();
    let loader = HistoricalDataLoader::new(config)
        .await
        .expect("Failed to create data loader");

    // Execute
    let (training_data, validation_data) = loader
        .load_training_data()
        .await
        .expect("Failed to load training data");

    // Verify
    assert!(!training_data.is_empty(), "Training data should not be empty");
    assert!(!validation_data.is_empty(), "Validation data should not be empty");

    // Verify split ratio (approximately 80/20)
    let total = training_data.len() + validation_data.len();
    let train_ratio = training_data.len() as f64 / total as f64;
    assert!(
        (train_ratio - 0.8).abs() < 0.1,
        "Train split ratio should be approximately 0.8, got {}",
        train_ratio
    );

    // Verify features structure
    let (features, targets) = &training_data[0];
    assert!(!features.prices.is_empty(), "Prices should not be empty");
    assert!(!features.volumes.is_empty(), "Volumes should not be empty");
    assert!(!features.technical_indicators.is_empty(), "Technical indicators should not be empty");
    assert!(!targets.is_empty(), "Targets should not be empty");

    println!("✅ Test passed: Loaded {} training samples, {} validation samples",
             training_data.len(), validation_data.len());
}

#[tokio::test]
#[ignore] // Requires test database setup
async fn test_time_range_filtering() {
    // Setup
    let pool = create_test_pool().await.expect("Failed to create test pool");
    setup_test_data(&pool).await.expect("Failed to setup test data");

    let mut config = create_test_config();
    config.time_range.start = Some(Utc::now() - chrono::Duration::minutes(30));
    config.time_range.end = Some(Utc::now());

    let loader = HistoricalDataLoader::new(config)
        .await
        .expect("Failed to create data loader");

    // Execute
    let (training_data, validation_data) = loader
        .load_training_data()
        .await
        .expect("Failed to load training data");

    // Verify data is within time range
    let total = training_data.len() + validation_data.len();
    assert!(
        total <= 30,
        "Should have at most 30 samples (30 minutes of data), got {}",
        total
    );

    println!("✅ Test passed: Time range filtering works correctly");
}

#[tokio::test]
#[ignore] // Requires test database setup
async fn test_symbol_filtering() {
    // Setup
    let pool = create_test_pool().await.expect("Failed to create test pool");
    setup_test_data(&pool).await.expect("Failed to setup test data");

    let mut config = create_test_config();
    config.symbols = vec!["TEST_SYMBOL".to_string()];

    let loader = HistoricalDataLoader::new(config)
        .await
        .expect("Failed to create data loader");

    // Execute
    let (training_data, _) = loader
        .load_training_data()
        .await
        .expect("Failed to load training data");

    // Verify all features are for TEST_SYMBOL
    for (features, _) in &training_data {
        // Note: We don't store symbol in FinancialFeatures, but we can verify
        // the data came from our test setup
        assert!(!features.prices.is_empty());
    }

    println!("✅ Test passed: Symbol filtering works correctly");
}

#[tokio::test]
#[ignore] // Requires test database setup
async fn test_data_validation() {
    // Setup
    let pool = create_test_pool().await.expect("Failed to create test pool");
    setup_test_data(&pool).await.expect("Failed to setup test data");

    let mut config = create_test_config();
    config.validation.min_samples = 1000; // Set unrealistically high

    let loader = HistoricalDataLoader::new(config)
        .await
        .expect("Failed to create data loader");

    // Execute - should fail due to insufficient samples
    let result = loader.load_training_data().await;

    // Verify
    assert!(
        result.is_err(),
        "Should fail with insufficient samples error"
    );

    let error_msg = result.unwrap_err().to_string();
    assert!(
        error_msg.contains("Insufficient data"),
        "Error should mention insufficient data, got: {}",
        error_msg
    );

    println!("✅ Test passed: Data validation rejects insufficient samples");
}

#[tokio::test]
#[ignore] // Requires test database setup
async fn test_feature_extraction() {
    // Setup
    let pool = create_test_pool().await.expect("Failed to create test pool");
    setup_test_data(&pool).await.expect("Failed to setup test data");

    let config = create_test_config();
    let loader = HistoricalDataLoader::new(config)
        .await
        .expect("Failed to create data loader");

    // Execute
    let (training_data, _) = loader
        .load_training_data()
        .await
        .expect("Failed to load training data");

    // Verify feature extraction
    let (features, _) = &training_data[0];

    // Check technical indicators
    assert!(
        features.technical_indicators.contains_key("spread_bps"),
        "Should have spread_bps indicator"
    );
    assert!(
        features.technical_indicators.contains_key("imbalance"),
        "Should have imbalance indicator"
    );

    // Check microstructure features
    assert!(features.microstructure.spread_bps > 0, "Spread should be positive");
    assert!(
        features.microstructure.imbalance.abs() <= 1.0,
        "Imbalance should be between -1 and 1"
    );

    // Check risk metrics
    assert!(
        features.risk_metrics.sharpe_ratio >= 0.0,
        "Sharpe ratio should be non-negative"
    );

    println!("✅ Test passed: Feature extraction produces valid features");
}