foxhunt/services/trading_service/tests/hot_swap_automation_tests.rs

//! Hot-Swap Automation Tests (TDD Approach)
//!
//! This test suite defines the expected behavior for automatic hot-swapping
//! of trained ML models into production ensemble. Tests written FIRST to drive
//! implementation.
//!
//! ## Test Coverage
//! 1. Automatic staging on training completion
//! 2. Validation (latency P99 < 200μs)
//! 3. Atomic swap (<1μs)
//! 4. Canary monitoring (5 minutes)
//! 5. Automatic rollback on failure
//! 6. Integration with HotSwapManager

#![allow(clippy::type_complexity, clippy::field_reassign_with_default)]

use std::sync::Arc;
use std::time::Duration;
use tokio::time::sleep;

use ml::ensemble::{CheckpointModel, CheckpointValidator, HotSwapManager, RollbackPolicy};
use ml::{Features, MLResult, ModelPrediction};
use trading_service::hot_swap_automation::{
    CanaryStatus, HotSwapAutomation, HotSwapConfig, TrainingEvent, ValidationStatus,
};

/// Helper: Create mock prediction function
fn create_mock_prediction_fn() -> Arc<dyn Fn(&Features) -> MLResult<ModelPrediction> + Send + Sync>
{
    Arc::new(|features: &Features| {
        let value = features.values.iter().sum::<f64>() / features.values.len() as f64;
        Ok(ModelPrediction::new("test".to_string(), value.tanh(), 0.85))
    })
}

/// Helper: Create slow prediction function (for validation failure)
fn create_slow_prediction_fn() -> Arc<dyn Fn(&Features) -> MLResult<ModelPrediction> + Send + Sync>
{
    Arc::new(|features: &Features| {
        std::thread::sleep(Duration::from_micros(100)); // 100μs > 50μs threshold
        let value = features.values.iter().sum::<f64>() / features.values.len() as f64;
        Ok(ModelPrediction::new("slow".to_string(), value.tanh(), 0.85))
    })
}

#[tokio::test]
async fn test_automatic_staging_on_training_complete() {
    // GIVEN: Hot-swap automation is running
    let hot_swap_manager = Arc::new(HotSwapManager::new(
        CheckpointValidator::new(),
        RollbackPolicy::default(),
    ));

    let config = HotSwapConfig::default();
    let automation = Arc::new(HotSwapAutomation::new(hot_swap_manager.clone(), config));

    // Register initial model
    let initial_model = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "checkpoint_v1.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));
    hot_swap_manager
        .register_model("DQN".to_string(), initial_model)
        .await
        .unwrap();

    // WHEN: Training completes with new checkpoint
    let new_checkpoint = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "checkpoint_v2.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));

    let event = TrainingEvent::new(
        "DQN".to_string(),
        "checkpoint_v2.safetensors".to_string(),
        new_checkpoint.clone(),
    );

    automation.handle_training_complete(event).await.unwrap();

    // THEN: Checkpoint should be staged automatically
    let status = automation.get_status("DQN").await.unwrap();
    assert_eq!(status.current_stage, "validated");
    assert_eq!(status.checkpoint_path, "checkpoint_v2.safetensors");
}

#[tokio::test]
async fn test_validation_latency_check() {
    // GIVEN: Hot-swap automation with strict validator
    let validator = CheckpointValidator::with_config(
        200,  // 200μs P99 threshold (realistic for CPU inference)
        1000, // 1000 test predictions
        (-1.0, 1.0),
    );

    let hot_swap_manager = Arc::new(HotSwapManager::new(validator, RollbackPolicy::default()));

    let config = HotSwapConfig::default();
    let automation = Arc::new(HotSwapAutomation::new(hot_swap_manager.clone(), config));

    // Register initial model
    let initial_model = Arc::new(CheckpointModel::new(
        "PPO".to_string(),
        "checkpoint_v1.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));
    hot_swap_manager
        .register_model("PPO".to_string(), initial_model)
        .await
        .unwrap();

    // WHEN: Training completes with fast checkpoint
    let fast_checkpoint = Arc::new(CheckpointModel::new(
        "PPO".to_string(),
        "checkpoint_fast.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));

    let event = TrainingEvent::new(
        "PPO".to_string(),
        "checkpoint_fast.safetensors".to_string(),
        fast_checkpoint,
    );

    automation.handle_training_complete(event).await.unwrap();

    // THEN: Validation should pass
    let status = automation.get_status("PPO").await.unwrap();
    assert!(matches!(
        status.validation_status,
        ValidationStatus::Passed { .. }
    ));
}

#[tokio::test]
async fn test_validation_rejects_slow_checkpoint() {
    // GIVEN: Hot-swap automation with strict validator
    let validator = CheckpointValidator::with_config(
        200,  // 200μs P99 threshold (realistic for CPU inference)
        1000, // 1000 test predictions
        (-1.0, 1.0),
    );

    let hot_swap_manager = Arc::new(HotSwapManager::new(validator, RollbackPolicy::default()));

    let config = HotSwapConfig::default();
    let automation = Arc::new(HotSwapAutomation::new(hot_swap_manager.clone(), config));

    // Register initial model
    let initial_model = Arc::new(CheckpointModel::new(
        "MAMBA2".to_string(),
        "checkpoint_v1.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));
    hot_swap_manager
        .register_model("MAMBA2".to_string(), initial_model)
        .await
        .unwrap();

    // WHEN: Training completes with slow checkpoint
    let slow_checkpoint = Arc::new(CheckpointModel::new(
        "MAMBA2".to_string(),
        "checkpoint_slow.safetensors".to_string(),
        create_slow_prediction_fn(),
    ));

    let event = TrainingEvent::new(
        "MAMBA2".to_string(),
        "checkpoint_slow.safetensors".to_string(),
        slow_checkpoint,
    );

    automation.handle_training_complete(event).await.unwrap();

    // THEN: Validation should fail and rollback
    let status = automation.get_status("MAMBA2").await.unwrap();
    assert!(matches!(
        status.validation_status,
        ValidationStatus::Failed { .. }
    ));
    assert_eq!(status.current_stage, "validation_failed");
}

#[tokio::test]
async fn test_atomic_swap_latency() {
    // GIVEN: Hot-swap automation is ready
    let hot_swap_manager = Arc::new(HotSwapManager::new(
        CheckpointValidator::new(),
        RollbackPolicy::default(),
    ));

    let config = HotSwapConfig::default();
    let automation = Arc::new(HotSwapAutomation::new(hot_swap_manager.clone(), config));

    // Register and stage checkpoint
    let initial_model = Arc::new(CheckpointModel::new(
        "TFT".to_string(),
        "checkpoint_v1.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));
    hot_swap_manager
        .register_model("TFT".to_string(), initial_model)
        .await
        .unwrap();

    let new_checkpoint = Arc::new(CheckpointModel::new(
        "TFT".to_string(),
        "checkpoint_v2.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));

    let event = TrainingEvent::new(
        "TFT".to_string(),
        "checkpoint_v2.safetensors".to_string(),
        new_checkpoint,
    );

    automation.handle_training_complete(event).await.unwrap();

    // Wait for validation
    sleep(Duration::from_millis(100)).await;

    // WHEN: Atomic swap is executed
    let swap_result = automation.execute_atomic_swap("TFT").await.unwrap();

    // THEN: Swap latency should be <100μs (testing threshold, production <1μs)
    assert!(
        swap_result.swap_latency_us < 100,
        "Swap latency {}μs exceeds 100μs",
        swap_result.swap_latency_us
    );
}

#[tokio::test]
async fn test_canary_monitoring_starts_after_swap() {
    // GIVEN: Hot-swap automation with short canary period
    let hot_swap_manager = Arc::new(HotSwapManager::new(
        CheckpointValidator::new(),
        RollbackPolicy::default(),
    ));

    let mut config = HotSwapConfig::default();
    config.canary_duration_secs = 1; // 1 second for testing

    let automation = Arc::new(HotSwapAutomation::new(hot_swap_manager.clone(), config));

    // Register and complete swap
    let initial_model = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "checkpoint_v1.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));
    hot_swap_manager
        .register_model("DQN".to_string(), initial_model)
        .await
        .unwrap();

    let new_checkpoint = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "checkpoint_v2.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));

    let event = TrainingEvent::new(
        "DQN".to_string(),
        "checkpoint_v2.safetensors".to_string(),
        new_checkpoint,
    );

    automation.handle_training_complete(event).await.unwrap();
    sleep(Duration::from_millis(100)).await;
    automation.execute_atomic_swap("DQN").await.unwrap();

    // WHEN: Checking canary status immediately after swap
    let status = automation.get_status("DQN").await.unwrap();

    // THEN: Canary monitoring should be active
    assert_eq!(status.current_stage, "canary_monitoring");
    assert!(matches!(
        status.canary_status,
        CanaryStatus::InProgress { .. }
    ));
}

#[tokio::test]
async fn test_canary_passes_and_completes() {
    // GIVEN: Hot-swap automation with very short canary period
    let hot_swap_manager = Arc::new(HotSwapManager::new(
        CheckpointValidator::new(),
        RollbackPolicy::default(),
    ));

    let mut config = HotSwapConfig::default();
    config.canary_duration_secs = 1; // 1 second for testing

    let automation = Arc::new(HotSwapAutomation::new(hot_swap_manager.clone(), config));

    // Complete full workflow
    let initial_model = Arc::new(CheckpointModel::new(
        "PPO".to_string(),
        "checkpoint_v1.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));
    hot_swap_manager
        .register_model("PPO".to_string(), initial_model)
        .await
        .unwrap();

    let new_checkpoint = Arc::new(CheckpointModel::new(
        "PPO".to_string(),
        "checkpoint_v2.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));

    let event = TrainingEvent::new(
        "PPO".to_string(),
        "checkpoint_v2.safetensors".to_string(),
        new_checkpoint,
    );

    automation.handle_training_complete(event).await.unwrap();
    sleep(Duration::from_millis(100)).await;
    automation.execute_atomic_swap("PPO").await.unwrap();

    // WHEN: Canary period completes (wait 2 seconds to be safe)
    sleep(Duration::from_secs(2)).await;

    // THEN: Canary should pass and workflow complete
    let status = automation.get_status("PPO").await.unwrap();
    assert!(matches!(status.canary_status, CanaryStatus::Passed));
    assert_eq!(status.current_stage, "completed");
}

#[tokio::test]
async fn test_automatic_rollback_on_canary_failure() {
    // GIVEN: Hot-swap automation with canary that will fail
    let hot_swap_manager = Arc::new(HotSwapManager::new(
        CheckpointValidator::new(),
        RollbackPolicy::strict(), // Strict policy for easier failure
    ));

    let mut config = HotSwapConfig::default();
    config.canary_duration_secs = 1;
    config.enable_automatic_rollback = true;

    let automation = Arc::new(HotSwapAutomation::new(hot_swap_manager.clone(), config));

    // Note: In real scenario, we'd inject failing metrics
    // For now, we test the rollback mechanism exists

    let initial_model = Arc::new(CheckpointModel::new(
        "MAMBA2".to_string(),
        "checkpoint_v1.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));
    hot_swap_manager
        .register_model("MAMBA2".to_string(), initial_model.clone())
        .await
        .unwrap();

    // Manually stage and commit swap
    let new_checkpoint = Arc::new(CheckpointModel::new(
        "MAMBA2".to_string(),
        "checkpoint_v2_bad.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));

    hot_swap_manager
        .stage_checkpoint("MAMBA2", new_checkpoint)
        .await
        .unwrap();
    hot_swap_manager.commit_swap("MAMBA2").await.unwrap();

    // WHEN: Rollback is triggered
    let rollback_result = automation.trigger_rollback("MAMBA2", "Test rollback").await;

    // THEN: Rollback should succeed and revert to previous checkpoint
    assert!(rollback_result.is_ok());

    let active = hot_swap_manager
        .get_active_checkpoint("MAMBA2")
        .await
        .unwrap();
    assert_eq!(active.checkpoint_path, "checkpoint_v1.safetensors");
}

#[tokio::test]
async fn test_concurrent_hot_swaps_for_different_models() {
    // GIVEN: Hot-swap automation with multiple models
    let hot_swap_manager = Arc::new(HotSwapManager::new(
        CheckpointValidator::new(),
        RollbackPolicy::default(),
    ));

    let config = HotSwapConfig::default();
    let automation = Arc::new(HotSwapAutomation::new(hot_swap_manager.clone(), config));

    // Register multiple models
    let models = vec!["DQN", "PPO", "MAMBA2", "TFT"];
    for model_id in &models {
        let model = Arc::new(CheckpointModel::new(
            model_id.to_string(),
            format!("{}_v1.safetensors", model_id),
            create_mock_prediction_fn(),
        ));
        hot_swap_manager
            .register_model(model_id.to_string(), model)
            .await
            .unwrap();
    }

    // WHEN: Multiple training events arrive concurrently
    let mut handles = vec![];
    for model_id in &models {
        let automation_clone = automation.clone();
        let model_id_clone = model_id.to_string();

        let handle = tokio::spawn(async move {
            let checkpoint = Arc::new(CheckpointModel::new(
                model_id_clone.clone(),
                format!("{}_v2.safetensors", model_id_clone),
                create_mock_prediction_fn(),
            ));

            let event = TrainingEvent::new(
                model_id_clone.clone(),
                format!("{}_v2.safetensors", model_id_clone),
                checkpoint,
            );

            automation_clone.handle_training_complete(event).await
        });

        handles.push(handle);
    }

    // Wait for all to complete
    for handle in handles {
        handle.await.unwrap().unwrap();
    }

    // THEN: All models should be staged independently
    for model_id in &models {
        let status = automation.get_status(model_id).await.unwrap();
        assert_eq!(status.current_stage, "validated");
    }
}

#[tokio::test]
async fn test_hot_swap_status_tracking() {
    // GIVEN: Hot-swap automation
    let hot_swap_manager = Arc::new(HotSwapManager::new(
        CheckpointValidator::new(),
        RollbackPolicy::default(),
    ));

    let config = HotSwapConfig::default();
    let automation = Arc::new(HotSwapAutomation::new(hot_swap_manager.clone(), config));

    // WHEN: Querying status for non-existent model
    let result = automation.get_status("NonExistent").await;

    // THEN: Should return error
    assert!(result.is_err());

    // WHEN: Registering model and triggering hot-swap workflow
    let model = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "checkpoint_v1.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));
    hot_swap_manager
        .register_model("DQN".to_string(), model)
        .await
        .unwrap();

    // Create and handle training event to generate status
    let new_checkpoint = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "checkpoint_v2.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));

    let event = TrainingEvent::new(
        "DQN".to_string(),
        "checkpoint_v2.safetensors".to_string(),
        new_checkpoint,
    );

    automation.handle_training_complete(event).await.unwrap();

    // THEN: Status should be available after training event
    let status = automation.get_status("DQN").await;
    assert!(status.is_ok());
}

#[tokio::test]
async fn test_disable_automatic_rollback() {
    // GIVEN: Hot-swap automation with automatic rollback disabled
    let hot_swap_manager = Arc::new(HotSwapManager::new(
        CheckpointValidator::new(),
        RollbackPolicy::default(),
    ));

    let mut config = HotSwapConfig::default();
    config.enable_automatic_rollback = false;

    let automation = Arc::new(HotSwapAutomation::new(hot_swap_manager.clone(), config));

    let initial_model = Arc::new(CheckpointModel::new(
        "TFT".to_string(),
        "checkpoint_v1.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));
    hot_swap_manager
        .register_model("TFT".to_string(), initial_model)
        .await
        .unwrap();

    // WHEN: Manual rollback is triggered (should still work)
    let new_checkpoint = Arc::new(CheckpointModel::new(
        "TFT".to_string(),
        "checkpoint_v2.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));
    hot_swap_manager
        .stage_checkpoint("TFT", new_checkpoint)
        .await
        .unwrap();
    hot_swap_manager.commit_swap("TFT").await.unwrap();

    let rollback_result = automation.trigger_rollback("TFT", "Manual test").await;

    // THEN: Manual rollback should still work
    assert!(rollback_result.is_ok());
}

#[tokio::test]
async fn test_full_e2e_hot_swap_workflow() {
    // GIVEN: Complete hot-swap automation setup
    let hot_swap_manager = Arc::new(HotSwapManager::new(
        CheckpointValidator::new(),
        RollbackPolicy::default(),
    ));

    let mut config = HotSwapConfig::default();
    config.canary_duration_secs = 1; // Short for testing

    let automation = Arc::new(HotSwapAutomation::new(hot_swap_manager.clone(), config));

    // Step 1: Register initial model
    let initial_model = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "checkpoint_v1.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));
    hot_swap_manager
        .register_model("DQN".to_string(), initial_model)
        .await
        .unwrap();

    // Step 2: Training completes
    let new_checkpoint = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "checkpoint_v2.safetensors".to_string(),
        create_mock_prediction_fn(),
    ));

    let event = TrainingEvent::new(
        "DQN".to_string(),
        "checkpoint_v2.safetensors".to_string(),
        new_checkpoint,
    );

    automation.handle_training_complete(event).await.unwrap();

    // Step 3: Verify validated (synchronous staging+validation)
    let status = automation.get_status("DQN").await.unwrap();
    assert_eq!(status.current_stage, "validated");

    // Step 4: Execute atomic swap
    sleep(Duration::from_millis(100)).await;
    let swap_result = automation.execute_atomic_swap("DQN").await.unwrap();
    assert!(swap_result.swap_latency_us < 100);

    // Step 5: Verify canary monitoring
    let status = automation.get_status("DQN").await.unwrap();
    assert_eq!(status.current_stage, "canary_monitoring");

    // Step 6: Wait for canary to complete
    sleep(Duration::from_secs(2)).await;

    // Step 7: Verify completion
    let status = automation.get_status("DQN").await.unwrap();
    assert!(matches!(status.canary_status, CanaryStatus::Passed));
    assert_eq!(status.current_stage, "completed");

    // Step 8: Verify new checkpoint is active
    let active = hot_swap_manager.get_active_checkpoint("DQN").await.unwrap();
    assert_eq!(active.checkpoint_path, "checkpoint_v2.safetensors");
}