foxhunt/services/trading_service/tests/prediction_generation_loop_tests.rs

//! TDD Tests for Background ML Prediction Generation Loop
//!
//! Test Coverage:
//! 1. Background task starts and runs continuously
//! 2. Predictions generated every 60 seconds
//! 3. Errors don't crash the loop (resilience)
//! 4. Graceful shutdown on SIGTERM
//! 5. Multiple symbols handled correctly

use anyhow::Result;
use sqlx::PgPool;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::broadcast;
use tokio::time::sleep;
use trading_service::ensemble_coordinator::EnsembleCoordinator;
use trading_service::prediction_generation_loop::{
    PredictionGenerationLoop, PredictionLoopConfig,
};

/// Helper to create test database pool
async fn create_test_pool() -> Result<PgPool> {
    let database_url = std::env::var("DATABASE_URL")
        .unwrap_or_else(|_| "postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt".to_string());

    let pool = PgPool::connect(&database_url).await?;
    Ok(pool)
}

/// Helper to create test ensemble coordinator with loaded models
async fn create_test_coordinator() -> Result<Arc<EnsembleCoordinator>> {
    use ml::model_factory;

    let coordinator = Arc::new(EnsembleCoordinator::new());

    // Register loaded models
    let dqn_model = model_factory::create_dqn_wrapper_with_id("DQN".to_string())?;
    let ppo_model = model_factory::create_ppo_wrapper_with_id("PPO".to_string())?;
    let tft_model = model_factory::create_tft_wrapper_with_id("TFT".to_string())?;

    coordinator
        .register_loaded_model("DQN".to_string(), dqn_model, 0.33)
        .await?;
    coordinator
        .register_loaded_model("PPO".to_string(), ppo_model, 0.33)
        .await?;
    coordinator
        .register_loaded_model("TFT".to_string(), tft_model, 0.34)
        .await?;

    Ok(coordinator)
}

/// Helper to insert mock market data for testing
async fn insert_mock_market_data(pool: &PgPool, symbol: &str) -> Result<()> {
    // Insert 100 bars of synthetic market data
    let now = chrono::Utc::now();

    for i in 0..100 {
        let timestamp = now - chrono::Duration::seconds(i * 60); // 1-minute bars
        let base_price = 4500.0 + (i as f64 * 0.5);

        sqlx::query(
            r#"
            INSERT INTO ohlcv_bars (timestamp, symbol, open, high, low, close, volume)
            VALUES ($1, $2, $3, $4, $5, $6, $7)
            ON CONFLICT (timestamp, symbol) DO NOTHING
            "#,
        )
        .bind(timestamp)
        .bind(symbol)
        .bind(base_price)
        .bind(base_price + 2.0)
        .bind(base_price - 2.0)
        .bind(base_price + 1.0)
        .bind(1_000_000.0)
        .execute(pool)
        .await?;
    }

    Ok(())
}

/// Helper to count predictions in database
async fn count_predictions(pool: &PgPool, symbol: &str) -> Result<i64> {
    let count: (i64,) = sqlx::query_as(
        r#"
        SELECT COUNT(*) FROM ensemble_predictions
        WHERE symbol = $1
        "#,
    )
    .bind(symbol)
    .fetch_one(pool)
    .await?;

    Ok(count.0)
}

/// Helper to get latest prediction
async fn get_latest_prediction(
    pool: &PgPool,
    symbol: &str,
) -> Result<Option<PredictionRecord>> {
    let prediction = sqlx::query_as::<_, PredictionRecord>(
        r#"
        SELECT
            id,
            prediction_timestamp,
            symbol,
            ensemble_action,
            ensemble_confidence,
            disagreement_rate,
            dqn_signal,
            ppo_signal,
            tft_signal
        FROM ensemble_predictions
        WHERE symbol = $1
        ORDER BY prediction_timestamp DESC
        LIMIT 1
        "#,
    )
    .bind(symbol)
    .fetch_optional(pool)
    .await?;

    Ok(prediction)
}

#[derive(Debug, sqlx::FromRow)]
struct PredictionRecord {
    id: uuid::Uuid,
    prediction_timestamp: chrono::DateTime<chrono::Utc>,
    symbol: String,
    ensemble_action: String,
    ensemble_confidence: f64,
    disagreement_rate: f64,
    dqn_signal: Option<f64>,
    ppo_signal: Option<f64>,
    tft_signal: Option<f64>,
}

// ================================================================================================
// TEST 1: Background task starts and runs continuously
// ================================================================================================

#[tokio::test]
#[ignore = "requires database and ML models"]
async fn test_background_task_starts_and_runs() -> Result<()> {
    let pool = create_test_pool().await?;
    let coordinator = create_test_coordinator().await?;

    // Create test config with short interval for testing
    let config = PredictionLoopConfig {
        prediction_interval: Duration::from_secs(2), // 2 seconds for testing
        symbols: vec!["ES.FUT".to_string()],
        account_id: "test_account".to_string(),
        strategy_id: "test_strategy".to_string(),
        node_id: "test_node".to_string(),
    };

    // Insert mock market data
    insert_mock_market_data(&pool, "ES.FUT").await?;

    // Create prediction loop
    let prediction_loop = PredictionGenerationLoop::new(coordinator, pool.clone(), config);

    // Create shutdown channel
    let (shutdown_tx, shutdown_rx) = broadcast::channel(1);

    // Spawn the loop in background
    let loop_handle = tokio::spawn(async move {
        prediction_loop.run(shutdown_rx).await
    });

    // Let it run for 5 seconds (should generate ~2 predictions)
    sleep(Duration::from_secs(5)).await;

    // Send shutdown signal
    shutdown_tx.send(()).ok();

    // Wait for loop to stop
    let result = loop_handle.await?;
    assert!(result.is_ok());

    // Verify predictions were generated
    let count = count_predictions(&pool, "ES.FUT").await?;
    assert!(count >= 2, "Expected at least 2 predictions, got {}", count);

    println!("✅ Test 1 passed: Background task started and ran continuously");
    Ok(())
}

// ================================================================================================
// TEST 2: Predictions generated every 60 seconds
// ================================================================================================

#[tokio::test]
#[ignore = "requires database and ML models"]
async fn test_predictions_generated_at_interval() -> Result<()> {
    let pool = create_test_pool().await?;
    let coordinator = create_test_coordinator().await?;

    // Create test config with 3-second interval for testing
    let config = PredictionLoopConfig {
        prediction_interval: Duration::from_secs(3),
        symbols: vec!["NQ.FUT".to_string()],
        account_id: "test_account".to_string(),
        strategy_id: "test_strategy".to_string(),
        node_id: "test_node".to_string(),
    };

    // Insert mock market data
    insert_mock_market_data(&pool, "NQ.FUT").await?;

    // Get initial count
    let initial_count = count_predictions(&pool, "NQ.FUT").await?;

    // Create prediction loop
    let prediction_loop = PredictionGenerationLoop::new(coordinator, pool.clone(), config);

    // Create shutdown channel
    let (shutdown_tx, shutdown_rx) = broadcast::channel(1);

    // Spawn the loop
    let loop_handle = tokio::spawn(async move {
        prediction_loop.run(shutdown_rx).await
    });

    // Wait for 3 intervals (9 seconds) - should generate 3 predictions
    sleep(Duration::from_secs(10)).await;

    // Send shutdown signal
    shutdown_tx.send(()).ok();
    loop_handle.await??;

    // Verify predictions were generated at the right frequency
    let final_count = count_predictions(&pool, "NQ.FUT").await?;
    let new_predictions = final_count - initial_count;

    assert!(
        new_predictions >= 3,
        "Expected at least 3 predictions, got {}",
        new_predictions
    );

    println!("✅ Test 2 passed: Predictions generated at correct interval");
    Ok(())
}

// ================================================================================================
// TEST 3: Errors don't crash the loop (resilience)
// ================================================================================================

#[tokio::test]
#[ignore = "requires database and ML models"]
async fn test_error_resilience() -> Result<()> {
    let pool = create_test_pool().await?;
    let coordinator = create_test_coordinator().await?;

    // Create test config with both valid and invalid symbols
    let config = PredictionLoopConfig {
        prediction_interval: Duration::from_secs(2),
        symbols: vec![
            "ES.FUT".to_string(),
            "INVALID_SYMBOL_NO_DATA".to_string(), // This will cause errors
            "NQ.FUT".to_string(),
        ],
        account_id: "test_account".to_string(),
        strategy_id: "test_strategy".to_string(),
        node_id: "test_node".to_string(),
    };

    // Insert market data for valid symbols only
    insert_mock_market_data(&pool, "ES.FUT").await?;
    insert_mock_market_data(&pool, "NQ.FUT").await?;

    // Create prediction loop
    let prediction_loop = PredictionGenerationLoop::new(coordinator, pool.clone(), config);

    // Create shutdown channel
    let (shutdown_tx, shutdown_rx) = broadcast::channel(1);

    // Spawn the loop
    let loop_handle = tokio::spawn(async move {
        prediction_loop.run(shutdown_rx).await
    });

    // Let it run for 5 seconds
    sleep(Duration::from_secs(5)).await;

    // Send shutdown signal
    shutdown_tx.send(()).ok();

    // Loop should NOT crash despite errors on invalid symbol
    let result = loop_handle.await?;
    assert!(result.is_ok(), "Loop crashed: {:?}", result);

    // Verify predictions were still generated for valid symbols
    let es_count = count_predictions(&pool, "ES.FUT").await?;
    let nq_count = count_predictions(&pool, "NQ.FUT").await?;

    assert!(es_count >= 1, "Expected predictions for ES.FUT, got {}", es_count);
    assert!(nq_count >= 1, "Expected predictions for NQ.FUT, got {}", nq_count);

    // Invalid symbol should have no predictions
    let invalid_count = count_predictions(&pool, "INVALID_SYMBOL_NO_DATA").await?;
    assert_eq!(invalid_count, 0);

    println!("✅ Test 3 passed: Loop resilient to errors");
    Ok(())
}

// ================================================================================================
// TEST 4: Graceful shutdown on SIGTERM
// ================================================================================================

#[tokio::test]
#[ignore = "requires database and ML models"]
async fn test_graceful_shutdown() -> Result<()> {
    let pool = create_test_pool().await?;
    let coordinator = create_test_coordinator().await?;

    let config = PredictionLoopConfig {
        prediction_interval: Duration::from_secs(2),
        symbols: vec!["ZN.FUT".to_string()],
        account_id: "test_account".to_string(),
        strategy_id: "test_strategy".to_string(),
        node_id: "test_node".to_string(),
    };

    // Insert mock market data
    insert_mock_market_data(&pool, "ZN.FUT").await?;

    // Create prediction loop
    let prediction_loop = PredictionGenerationLoop::new(coordinator, pool.clone(), config);

    // Create shutdown channel
    let (shutdown_tx, shutdown_rx) = broadcast::channel(1);

    // Spawn the loop
    let loop_handle = tokio::spawn(async move {
        prediction_loop.run(shutdown_rx).await
    });

    // Let it run for a bit
    sleep(Duration::from_secs(3)).await;

    // Send shutdown signal
    let send_result = shutdown_tx.send(());
    assert!(send_result.is_ok(), "Failed to send shutdown signal");

    // Wait for loop to stop (should be immediate)
    let result = tokio::time::timeout(Duration::from_secs(5), loop_handle).await;

    assert!(result.is_ok(), "Loop did not shutdown within 5 seconds");
    assert!(result.unwrap().is_ok(), "Loop failed during shutdown");

    println!("✅ Test 4 passed: Graceful shutdown successful");
    Ok(())
}

// ================================================================================================
// TEST 5: Multiple symbols handled correctly
// ================================================================================================

#[tokio::test]
#[ignore = "requires database and ML models"]
async fn test_multiple_symbols() -> Result<()> {
    let pool = create_test_pool().await?;
    let coordinator = create_test_coordinator().await?;

    // Create test config with 4 symbols (default production config)
    let config = PredictionLoopConfig {
        prediction_interval: Duration::from_secs(2),
        symbols: vec![
            "ES.FUT".to_string(),
            "NQ.FUT".to_string(),
            "ZN.FUT".to_string(),
            "6E.FUT".to_string(),
        ],
        account_id: "test_account".to_string(),
        strategy_id: "test_strategy".to_string(),
        node_id: "test_node".to_string(),
    };

    // Insert mock market data for all symbols
    for symbol in &config.symbols {
        insert_mock_market_data(&pool, symbol).await?;
    }

    // Create prediction loop
    let prediction_loop = PredictionGenerationLoop::new(coordinator, pool.clone(), config.clone());

    // Create shutdown channel
    let (shutdown_tx, shutdown_rx) = broadcast::channel(1);

    // Spawn the loop
    let loop_handle = tokio::spawn(async move {
        prediction_loop.run(shutdown_rx).await
    });

    // Let it run for 5 seconds (should generate at least 2 cycles)
    sleep(Duration::from_secs(5)).await;

    // Send shutdown signal
    shutdown_tx.send(()).ok();
    loop_handle.await??;

    // Verify predictions were generated for ALL symbols
    for symbol in &config.symbols {
        let count = count_predictions(&pool, symbol).await?;
        assert!(
            count >= 2,
            "Expected at least 2 predictions for {}, got {}",
            symbol,
            count
        );

        // Verify prediction structure
        if let Some(prediction) = get_latest_prediction(&pool, symbol).await? {
            assert_eq!(prediction.symbol, *symbol);
            assert!(prediction.ensemble_confidence >= 0.0 && prediction.ensemble_confidence <= 1.0);
            assert!(prediction.disagreement_rate >= 0.0 && prediction.disagreement_rate <= 1.0);
            assert!(
                vec!["BUY", "SELL", "HOLD"].contains(&prediction.ensemble_action.as_str())
            );

            // Verify per-model signals are present
            assert!(prediction.dqn_signal.is_some());
            assert!(prediction.ppo_signal.is_some());
            assert!(prediction.tft_signal.is_some());

            println!("✅ Verified prediction for {}: action={}, confidence={:.3}",
                symbol, prediction.ensemble_action, prediction.ensemble_confidence);
        }
    }

    println!("✅ Test 5 passed: Multiple symbols handled correctly");
    Ok(())
}

// ================================================================================================
// TEST 6: Verify prediction data structure
// ================================================================================================

#[tokio::test]
#[ignore = "requires database and ML models"]
async fn test_prediction_data_structure() -> Result<()> {
    let pool = create_test_pool().await?;
    let coordinator = create_test_coordinator().await?;

    let config = PredictionLoopConfig {
        prediction_interval: Duration::from_secs(2),
        symbols: vec!["ES.FUT".to_string()],
        account_id: "test_account_123".to_string(),
        strategy_id: "test_strategy_v1".to_string(),
        node_id: "test_node_001".to_string(),
    };

    // Insert mock market data
    insert_mock_market_data(&pool, "ES.FUT").await?;

    // Create prediction loop
    let prediction_loop = PredictionGenerationLoop::new(coordinator, pool.clone(), config.clone());

    // Create shutdown channel
    let (shutdown_tx, shutdown_rx) = broadcast::channel(1);

    // Spawn the loop
    let loop_handle = tokio::spawn(async move {
        prediction_loop.run(shutdown_rx).await
    });

    // Let it generate one prediction
    sleep(Duration::from_secs(3)).await;

    // Send shutdown signal
    shutdown_tx.send(()).ok();
    loop_handle.await??;

    // Verify prediction structure in detail
    let prediction = sqlx::query(
        r#"
        SELECT
            id,
            prediction_timestamp,
            symbol,
            account_id,
            strategy_id,
            ensemble_action,
            ensemble_signal,
            ensemble_confidence,
            disagreement_rate,
            dqn_signal,
            dqn_confidence,
            dqn_weight,
            dqn_vote,
            ppo_signal,
            ppo_confidence,
            ppo_weight,
            ppo_vote,
            tft_signal,
            tft_confidence,
            tft_weight,
            tft_vote,
            node_id,
            inference_latency_us
        FROM ensemble_predictions
        WHERE symbol = $1
        ORDER BY prediction_timestamp DESC
        LIMIT 1
        "#,
    )
    .bind("ES.FUT")
    .fetch_one(&pool)
    .await?;

    // Verify all fields are populated correctly
    let account_id: String = prediction.get("account_id");
    let strategy_id: String = prediction.get("strategy_id");
    let node_id: String = prediction.get("node_id");

    assert_eq!(account_id, config.account_id);
    assert_eq!(strategy_id, config.strategy_id);
    assert_eq!(node_id, config.node_id);

    // Verify ensemble fields
    let ensemble_signal: f64 = prediction.get("ensemble_signal");
    let ensemble_confidence: f64 = prediction.get("ensemble_confidence");
    let disagreement_rate: f64 = prediction.get("disagreement_rate");

    assert!(ensemble_signal >= -1.0 && ensemble_signal <= 1.0);
    assert!(ensemble_confidence >= 0.0 && ensemble_confidence <= 1.0);
    assert!(disagreement_rate >= 0.0 && disagreement_rate <= 1.0);

    // Verify per-model fields are present
    let dqn_signal: Option<f64> = prediction.get("dqn_signal");
    let dqn_confidence: Option<f64> = prediction.get("dqn_confidence");
    let dqn_weight: Option<f64> = prediction.get("dqn_weight");
    let dqn_vote: Option<String> = prediction.get("dqn_vote");

    assert!(dqn_signal.is_some());
    assert!(dqn_confidence.is_some());
    assert!(dqn_weight.is_some());
    assert!(dqn_vote.is_some());

    // Verify latency is reasonable
    let inference_latency_us: i32 = prediction.get("inference_latency_us");
    assert!(inference_latency_us > 0);
    assert!(inference_latency_us < 1_000_000); // Should be < 1 second

    println!("✅ Test 6 passed: Prediction data structure verified");
    Ok(())
}