foxhunt/services/trading_service/tests/ensemble_risk_integration_test.rs

//! Integration tests for Ensemble Risk Manager
//!
//! Tests all risk management scenarios:
//! - Low confidence rejection
//! - High disagreement rejection
//! - Per-model circuit breakers
//! - Cascade failure detection
//! - VaR integration
//! - Model cooldown and recovery

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio;

use ml::ensemble::{EnsembleDecision, ModelVote, TradingAction};
use trading_service::ensemble_risk_manager::{
    EnsembleRiskConfig, EnsembleRiskManager, RiskValidationResult,
};

#[tokio::test]
async fn test_low_confidence_rejection() {
    let config = EnsembleRiskConfig {
        min_confidence_threshold: 0.70,
        ..Default::default()
    };
    let manager = EnsembleRiskManager::new(config);

    // Create low confidence decision
    let decision = EnsembleDecision::new(
        TradingAction::Buy,
        0.55, // Below 70% threshold
        0.60,
        0.20,
        HashMap::new(),
    );

    let result = manager
        .validate_prediction(&decision, "TEST_ACCOUNT")
        .await
        .unwrap();

    assert!(!result.approved, "Low confidence should be rejected");
    assert!(
        result.rejection_reason.is_some(),
        "Rejection reason should be provided"
    );
    assert!(
        result
            .rejection_reason
            .unwrap()
            .contains("Low confidence"),
        "Rejection reason should mention confidence"
    );
    assert_eq!(result.confidence, 0.55);
}

#[tokio::test]
async fn test_high_disagreement_rejection() {
    let config = EnsembleRiskConfig {
        max_disagreement_rate: 0.40,
        ..Default::default()
    };
    let manager = EnsembleRiskManager::new(config);

    // Create high disagreement decision
    let decision = EnsembleDecision::new(
        TradingAction::Buy,
        0.75,
        0.65,
        0.55, // Above 40% disagreement threshold
        HashMap::new(),
    );

    let result = manager
        .validate_prediction(&decision, "TEST_ACCOUNT")
        .await
        .unwrap();

    assert!(!result.approved, "High disagreement should be rejected");
    assert!(
        result
            .rejection_reason
            .unwrap()
            .contains("High disagreement"),
        "Rejection reason should mention disagreement"
    );
}

#[tokio::test]
async fn test_model_circuit_breaker_after_consecutive_errors() {
    let config = EnsembleRiskConfig {
        max_consecutive_errors: 3,
        ..Default::default()
    };
    let manager = EnsembleRiskManager::new(config);

    // Register model
    manager.register_model("DQN".to_string()).await.unwrap();

    // Verify model starts enabled
    let health = manager.get_model_health("DQN").await.unwrap();
    assert!(health.enabled, "Model should start enabled");
    assert_eq!(
        health.consecutive_errors, 0,
        "Should start with 0 errors"
    );

    // Record 3 consecutive errors
    for i in 1..=3 {
        manager
            .record_prediction_result("DQN", false)
            .await
            .unwrap();

        let health = manager.get_model_health("DQN").await.unwrap();
        assert_eq!(health.consecutive_errors, i);
        assert_eq!(health.failed_predictions, i as u64);
    }

    // Verify model is disabled after 3 errors
    let health = manager.get_model_health("DQN").await.unwrap();
    assert!(!health.enabled, "Model should be disabled after 3 errors");
    assert_eq!(health.consecutive_errors, 3);
    assert!(
        health.disabled_at.is_some(),
        "Disabled timestamp should be set"
    );
    assert!(
        health.cooldown_until.is_some(),
        "Cooldown period should be set"
    );
}

#[tokio::test]
async fn test_cascade_failure_detection() {
    let config = EnsembleRiskConfig {
        max_consecutive_errors: 2,
        cascade_failure_threshold: 2, // 2+ models = cascade
        cascade_detection_window_secs: 60,
        ..Default::default()
    };
    let manager = EnsembleRiskManager::new(config);

    // Register 3 models
    manager.register_model("DQN".to_string()).await.unwrap();
    manager.register_model("PPO".to_string()).await.unwrap();
    manager.register_model("TFT".to_string()).await.unwrap();

    // Verify ensemble is operational
    assert!(
        manager.is_operational().await,
        "Ensemble should start operational"
    );

    // Fail first model (2 errors to disable)
    for _ in 0..2 {
        manager
            .record_prediction_result("DQN", false)
            .await
            .unwrap();
    }

    let dqn_health = manager.get_model_health("DQN").await.unwrap();
    assert!(!dqn_health.enabled, "DQN should be disabled");

    // Ensemble should still be operational (only 1 model failed)
    assert!(
        manager.is_operational().await,
        "Ensemble should be operational with 1 failed model"
    );

    // Fail second model to trigger cascade
    for _ in 0..2 {
        manager
            .record_prediction_result("PPO", false)
            .await
            .unwrap();
    }

    let ppo_health = manager.get_model_health("PPO").await.unwrap();
    assert!(!ppo_health.enabled, "PPO should be disabled");

    // Verify cascade detected
    let cascade_state = manager.get_cascade_state().await;
    assert!(
        cascade_state.is_cascading,
        "Cascade should be detected with 2 failed models"
    );
    assert_eq!(
        cascade_state.failed_models.len(),
        2,
        "Should track 2 failed models"
    );
    assert!(
        cascade_state.failed_models.contains(&"DQN".to_string()),
        "Should include DQN"
    );
    assert!(
        cascade_state.failed_models.contains(&"PPO".to_string()),
        "Should include PPO"
    );

    // Ensemble should NOT be operational
    assert!(
        !manager.is_operational().await,
        "Ensemble should not be operational during cascade"
    );

    // Predictions should be rejected during cascade
    let decision = EnsembleDecision::new(
        TradingAction::Buy,
        0.85,
        0.70,
        0.15,
        HashMap::new(),
    );

    let result = manager
        .validate_prediction(&decision, "TEST_ACCOUNT")
        .await
        .unwrap();

    assert!(
        !result.approved,
        "Predictions should be rejected during cascade"
    );
    assert!(result.cascade_detected || result.rejection_reason.is_some());
}

#[tokio::test]
async fn test_successful_predictions_reset_consecutive_errors() {
    let config = EnsembleRiskConfig {
        max_consecutive_errors: 3,
        ..Default::default()
    };
    let manager = EnsembleRiskManager::new(config);

    manager.register_model("DQN".to_string()).await.unwrap();

    // Record 2 errors
    manager
        .record_prediction_result("DQN", false)
        .await
        .unwrap();
    manager
        .record_prediction_result("DQN", false)
        .await
        .unwrap();

    let health = manager.get_model_health("DQN").await.unwrap();
    assert_eq!(health.consecutive_errors, 2);
    assert!(health.enabled, "Should still be enabled");

    // Record successful prediction
    manager
        .record_prediction_result("DQN", true)
        .await
        .unwrap();

    // Verify consecutive errors reset
    let health = manager.get_model_health("DQN").await.unwrap();
    assert_eq!(
        health.consecutive_errors, 0,
        "Consecutive errors should reset"
    );
    assert!(health.enabled, "Should remain enabled");
    assert_eq!(health.successful_predictions, 1);
    assert_eq!(health.failed_predictions, 2);
}

#[tokio::test]
async fn test_model_cooldown_and_recovery() {
    let config = EnsembleRiskConfig {
        max_consecutive_errors: 2,
        model_cooldown_period_secs: 1, // 1 second for fast testing
        ..Default::default()
    };
    let manager = EnsembleRiskManager::new(config);

    manager.register_model("DQN".to_string()).await.unwrap();

    // Disable model with 2 errors
    manager
        .record_prediction_result("DQN", false)
        .await
        .unwrap();
    manager
        .record_prediction_result("DQN", false)
        .await
        .unwrap();

    let health = manager.get_model_health("DQN").await.unwrap();
    assert!(!health.enabled, "Model should be disabled");
    assert!(health.is_in_cooldown(), "Should be in cooldown");

    // Try to enable immediately (should fail - still in cooldown)
    let enabled = manager.try_enable_model("DQN").await.unwrap();
    assert!(
        !enabled,
        "Should not enable during cooldown period"
    );

    let health = manager.get_model_health("DQN").await.unwrap();
    assert!(
        !health.enabled,
        "Model should still be disabled during cooldown"
    );

    // Wait for cooldown period to expire
    tokio::time::sleep(Duration::from_secs(2)).await;

    // Try to enable after cooldown
    let enabled = manager.try_enable_model("DQN").await.unwrap();
    assert!(enabled, "Should enable after cooldown period");

    let health = manager.get_model_health("DQN").await.unwrap();
    assert!(health.enabled, "Model should be enabled");
    assert_eq!(
        health.consecutive_errors, 0,
        "Errors should be reset"
    );
    assert!(!health.is_in_cooldown(), "Should not be in cooldown");
}

#[tokio::test]
async fn test_approved_prediction_with_good_metrics() {
    let config = EnsembleRiskConfig::default();
    let manager = EnsembleRiskManager::new(config);

    // Create high-quality decision
    let decision = EnsembleDecision::new(
        TradingAction::Buy,
        0.85, // High confidence
        0.80,
        0.10, // Low disagreement
        HashMap::new(),
    );

    let result = manager
        .validate_prediction(&decision, "TEST_ACCOUNT")
        .await
        .unwrap();

    assert!(result.approved, "Good prediction should be approved");
    assert!(
        result.rejection_reason.is_none(),
        "No rejection reason for approved prediction"
    );
    assert_eq!(result.confidence, 0.85);
    assert_eq!(result.disagreement_rate, 0.10);
    assert!(
        result.validation_latency_us > 0,
        "Should track validation latency"
    );
}

#[tokio::test]
async fn test_cascade_manual_reset() {
    let config = EnsembleRiskConfig {
        max_consecutive_errors: 1,
        cascade_failure_threshold: 2,
        ..Default::default()
    };
    let manager = EnsembleRiskManager::new(config);

    // Register and fail 2 models
    manager.register_model("DQN".to_string()).await.unwrap();
    manager.register_model("PPO".to_string()).await.unwrap();

    manager
        .record_prediction_result("DQN", false)
        .await
        .unwrap();
    manager
        .record_prediction_result("PPO", false)
        .await
        .unwrap();

    // Verify cascade detected
    let cascade_state = manager.get_cascade_state().await;
    assert!(cascade_state.is_cascading);

    // Manually reset cascade
    manager.reset_cascade_state().await.unwrap();

    // Verify cascade cleared
    let cascade_state = manager.get_cascade_state().await;
    assert!(
        !cascade_state.is_cascading,
        "Cascade should be cleared"
    );
    assert_eq!(
        cascade_state.failed_models.len(),
        0,
        "Failed models should be cleared"
    );
}

#[tokio::test]
async fn test_multiple_model_health_tracking() {
    let config = EnsembleRiskConfig {
        max_consecutive_errors: 3,
        ..Default::default()
    };
    let manager = EnsembleRiskManager::new(config);

    // Register 3 models
    manager.register_model("DQN".to_string()).await.unwrap();
    manager.register_model("PPO".to_string()).await.unwrap();
    manager.register_model("TFT".to_string()).await.unwrap();

    // Different success rates for each model
    // DQN: 2 success, 1 failure
    manager.record_prediction_result("DQN", true).await.unwrap();
    manager.record_prediction_result("DQN", true).await.unwrap();
    manager
        .record_prediction_result("DQN", false)
        .await
        .unwrap();

    // PPO: 1 success, 2 failures (but not consecutive)
    manager.record_prediction_result("PPO", true).await.unwrap();
    manager
        .record_prediction_result("PPO", false)
        .await
        .unwrap();
    manager
        .record_prediction_result("PPO", false)
        .await
        .unwrap();

    // TFT: 3 successes
    manager.record_prediction_result("TFT", true).await.unwrap();
    manager.record_prediction_result("TFT", true).await.unwrap();
    manager.record_prediction_result("TFT", true).await.unwrap();

    // Check all health statuses
    let all_health = manager.get_all_model_health().await;
    assert_eq!(all_health.len(), 3);

    let dqn_health = all_health.get("DQN").unwrap();
    assert_eq!(dqn_health.successful_predictions, 2);
    assert_eq!(dqn_health.failed_predictions, 1);
    assert_eq!(dqn_health.consecutive_errors, 1); // Last was failure
    assert!(dqn_health.enabled);

    let ppo_health = all_health.get("PPO").unwrap();
    assert_eq!(ppo_health.successful_predictions, 1);
    assert_eq!(ppo_health.failed_predictions, 2);
    assert_eq!(ppo_health.consecutive_errors, 2); // Last 2 were failures
    assert!(ppo_health.enabled); // Not yet at threshold

    let tft_health = all_health.get("TFT").unwrap();
    assert_eq!(tft_health.successful_predictions, 3);
    assert_eq!(tft_health.failed_predictions, 0);
    assert_eq!(tft_health.consecutive_errors, 0);
    assert!(tft_health.enabled);

    // Verify enabled count
    assert_eq!(
        manager.enabled_model_count().await,
        3,
        "All models should be enabled"
    );
}

#[tokio::test]
async fn test_cascade_detection_window_expiry() {
    let config = EnsembleRiskConfig {
        max_consecutive_errors: 1,
        cascade_failure_threshold: 2,
        cascade_detection_window_secs: 1, // 1 second window
        ..Default::default()
    };
    let manager = EnsembleRiskManager::new(config);

    manager.register_model("DQN".to_string()).await.unwrap();
    manager.register_model("PPO".to_string()).await.unwrap();

    // Fail first model
    manager
        .record_prediction_result("DQN", false)
        .await
        .unwrap();

    let cascade_state = manager.get_cascade_state().await;
    assert_eq!(cascade_state.failed_models.len(), 1);
    assert!(!cascade_state.is_cascading);

    // Wait for detection window to expire
    tokio::time::sleep(Duration::from_secs(2)).await;

    // Fail second model (should reset window, not trigger cascade)
    manager
        .record_prediction_result("PPO", false)
        .await
        .unwrap();

    let cascade_state = manager.get_cascade_state().await;
    // Window should have reset, so only PPO in failed list
    assert!(
        cascade_state.failed_models.len() <= 2,
        "Detection window should have reset"
    );
}

#[tokio::test]
async fn test_error_rate_calculation() {
    let config = EnsembleRiskConfig {
        max_consecutive_errors: 10, // High threshold to prevent disabling
        ..Default::default()
    };
    let manager = EnsembleRiskManager::new(config);

    manager.register_model("DQN".to_string()).await.unwrap();

    // Record 7 successes and 3 failures
    for _ in 0..7 {
        manager
            .record_prediction_result("DQN", true)
            .await
            .unwrap();
    }
    for _ in 0..3 {
        manager
            .record_prediction_result("DQN", false)
            .await
            .unwrap();
    }

    let health = manager.get_model_health("DQN").await.unwrap();
    assert_eq!(health.total_predictions, 10);
    assert_eq!(health.successful_predictions, 7);
    assert_eq!(health.failed_predictions, 3);

    let error_rate = health.error_rate();
    assert!(
        (error_rate - 0.30).abs() < 0.01,
        "Error rate should be 30%"
    );
}

#[tokio::test]
async fn test_validation_tracks_disabled_models() {
    let config = EnsembleRiskConfig {
        max_consecutive_errors: 2,
        ..Default::default()
    };
    let manager = EnsembleRiskManager::new(config);

    // Register 3 models
    manager.register_model("DQN".to_string()).await.unwrap();
    manager.register_model("PPO".to_string()).await.unwrap();
    manager.register_model("TFT".to_string()).await.unwrap();

    // Disable DQN
    manager
        .record_prediction_result("DQN", false)
        .await
        .unwrap();
    manager
        .record_prediction_result("DQN", false)
        .await
        .unwrap();

    // Make good prediction
    let decision = EnsembleDecision::new(
        TradingAction::Buy,
        0.85,
        0.80,
        0.15,
        HashMap::new(),
    );

    let result = manager
        .validate_prediction(&decision, "TEST_ACCOUNT")
        .await
        .unwrap();

    assert!(result.approved);
    assert_eq!(
        result.disabled_models.len(),
        1,
        "Should track 1 disabled model"
    );
    assert!(
        result.disabled_models.contains(&"DQN".to_string()),
        "Should list DQN as disabled"
    );
}