foxhunt/ml/tests/ensemble_hot_swap_test.rs

//! Integration tests for ensemble checkpoint hot-swapping
//!
//! Tests the complete hot-swap workflow:
//! 1. Load DQN checkpoint into active buffer
//! 2. Stage new checkpoint in shadow buffer
//! 3. Validate staged checkpoint (1000 predictions, latency checks)
//! 4. Commit atomic swap (<1μs)
//! 5. Verify active buffer now points to new checkpoint
//! 6. Test rollback mechanism

use ml::ensemble::{
    CheckpointModel, CheckpointValidator, HotSwapManager, RollbackPolicy,
    ValidationResult, CanaryResult, EnsembleMetrics,
};
use ml::{Features, ModelPrediction, MLResult};
use std::sync::Arc;
use std::time::Instant;

/// Mock prediction function for DQN epoch 30
fn create_dqn_epoch_30_predictor() -> Arc<dyn Fn(&Features) -> MLResult<ModelPrediction> + Send + Sync> {
    Arc::new(|features: &Features| {
        let value = (features.values.iter().sum::<f64>() / features.values.len() as f64) * 0.8;
        Ok(ModelPrediction::new("DQN_epoch_30".to_string(), value.tanh(), 0.78))
    })
}

/// Mock prediction function for DQN epoch 50 (improved)
fn create_dqn_epoch_50_predictor() -> Arc<dyn Fn(&Features) -> MLResult<ModelPrediction> + Send + Sync> {
    Arc::new(|features: &Features| {
        let value = (features.values.iter().sum::<f64>() / features.values.len() as f64) * 0.9;
        Ok(ModelPrediction::new("DQN_epoch_50".to_string(), value.tanh(), 0.85))
    })
}

#[tokio::test]
async fn test_hot_swap_workflow_complete() {
    println!("\n=== Hot-Swap Workflow Test ===\n");

    // Initialize hot-swap manager
    let validator = CheckpointValidator::new();
    let policy = RollbackPolicy::default();
    let manager = HotSwapManager::new(validator, policy);

    // Step 1: Load DQN epoch 30 into active buffer
    println!("Step 1: Loading DQN epoch 30 into active buffer...");
    let checkpoint_v30 = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "ml/checkpoints/dqn/checkpoint_epoch_30.safetensors".to_string(),
        create_dqn_epoch_30_predictor(),
    ));

    manager
        .register_model("DQN".to_string(), checkpoint_v30)
        .await
        .expect("Failed to register model");

    let active = manager.get_active_checkpoint("DQN").await.unwrap();
    assert_eq!(active.checkpoint_path, "ml/checkpoints/dqn/checkpoint_epoch_30.safetensors");
    println!("✓ Active checkpoint: {}", active.checkpoint_path);

    // Step 2: Stage DQN epoch 50 in shadow buffer
    println!("\nStep 2: Staging DQN epoch 50 in shadow buffer...");
    let checkpoint_v50 = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "ml/checkpoints/dqn/checkpoint_epoch_50.safetensors".to_string(),
        create_dqn_epoch_50_predictor(),
    ));

    manager
        .stage_checkpoint("DQN", checkpoint_v50)
        .await
        .expect("Failed to stage checkpoint");
    println!("✓ Staged checkpoint in shadow buffer");

    // Step 3: Validate staged checkpoint
    println!("\nStep 3: Validating staged checkpoint (1000 predictions)...");
    let validation_start = Instant::now();
    let validation = manager
        .validate_staged_checkpoint("DQN")
        .await
        .expect("Failed to validate checkpoint");
    let validation_duration = validation_start.elapsed();

    assert!(validation.passed, "Validation failed: {:?}", validation.failure_reason);
    assert!(validation.p99_latency_us < 50, "P99 latency {}μs exceeds 50μs", validation.p99_latency_us);
    assert!(validation.predictions_in_range >= 950, "Only {} predictions in range", validation.predictions_in_range);

    println!("✓ Validation PASSED:");
    println!("  - Avg latency: {}μs", validation.avg_latency_us);
    println!("  - P99 latency: {}μs", validation.p99_latency_us);
    println!("  - Predictions: {}/{} in range ({:.1}%)",
        validation.predictions_in_range,
        validation.predictions_validated,
        (validation.predictions_in_range as f64 / validation.predictions_validated as f64) * 100.0
    );
    println!("  - Validation duration: {}ms", validation_duration.as_millis());

    // Record metrics
    EnsembleMetrics::record_validation(
        "DQN",
        true,
        validation_duration.as_millis() as u64,
        validation.p99_latency_us,
    );

    // Step 4: Commit atomic swap
    println!("\nStep 4: Committing atomic swap...");
    let swap_latency = manager
        .commit_swap("DQN")
        .await
        .expect("Failed to commit swap");

    assert!(
        swap_latency.as_micros() < 100,
        "Swap latency {}μs exceeds 100μs",
        swap_latency.as_micros()
    );

    println!("✓ Atomic swap completed in {}μs", swap_latency.as_micros());

    // Record metrics
    EnsembleMetrics::record_swap_success("DQN", swap_latency.as_micros() as u64);

    // Step 5: Verify active checkpoint is now epoch 50
    println!("\nStep 5: Verifying active checkpoint...");
    let active = manager.get_active_checkpoint("DQN").await.unwrap();
    assert_eq!(active.checkpoint_path, "ml/checkpoints/dqn/checkpoint_epoch_50.safetensors");
    println!("✓ Active checkpoint: {}", active.checkpoint_path);

    // Test prediction with new checkpoint
    let features = Features::new(
        vec![0.5, 0.6, 0.7, 0.8, 0.9],
        vec!["f1".to_string(), "f2".to_string(), "f3".to_string(), "f4".to_string(), "f5".to_string()],
    );

    let prediction = active.predict(&features).unwrap();
    assert_eq!(prediction.model_id, "DQN_epoch_50");
    assert!(prediction.confidence >= 0.85);
    println!("✓ Prediction with new checkpoint: value={:.4}, confidence={:.4}", prediction.value, prediction.confidence);

    println!("\n=== Hot-Swap Workflow Test PASSED ===\n");
}

#[tokio::test]
async fn test_hot_swap_rollback() {
    println!("\n=== Hot-Swap Rollback Test ===\n");

    let validator = CheckpointValidator::new();
    let policy = RollbackPolicy::default();
    let manager = HotSwapManager::new(validator, policy);

    // Register initial checkpoint
    let checkpoint_v30 = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "checkpoint_epoch_30.safetensors".to_string(),
        create_dqn_epoch_30_predictor(),
    ));

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

    // Stage and commit new checkpoint
    let checkpoint_v50 = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "checkpoint_epoch_50.safetensors".to_string(),
        create_dqn_epoch_50_predictor(),
    ));

    manager.stage_checkpoint("DQN", checkpoint_v50).await.unwrap();
    manager.commit_swap("DQN").await.unwrap();

    // Verify new checkpoint is active
    let active = manager.get_active_checkpoint("DQN").await.unwrap();
    assert_eq!(active.checkpoint_path, "checkpoint_epoch_50.safetensors");
    println!("✓ Swapped to epoch 50");

    // Simulate failure and rollback
    println!("Simulating failure and rolling back...");
    manager.rollback("DQN").await.unwrap();

    // Verify rollback restored epoch 30
    let active = manager.get_active_checkpoint("DQN").await.unwrap();
    assert_eq!(active.checkpoint_path, "checkpoint_epoch_30.safetensors");
    println!("✓ Rolled back to epoch 30");

    // Record metrics
    EnsembleMetrics::record_swap_rollback("DQN");
    EnsembleMetrics::record_rollback("DQN", "test_failure");

    println!("\n=== Rollback Test PASSED ===\n");
}

#[tokio::test]
async fn test_swap_latency_benchmark() {
    println!("\n=== Swap Latency Benchmark ===\n");

    let validator = CheckpointValidator::new();
    let policy = RollbackPolicy::default();
    let manager = HotSwapManager::new(validator, policy);

    // Register initial checkpoint
    let checkpoint_v1 = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "checkpoint_v1.safetensors".to_string(),
        create_dqn_epoch_30_predictor(),
    ));

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

    // Run 100 swap operations to benchmark
    let mut swap_latencies = Vec::new();

    for i in 0..100 {
        let checkpoint = Arc::new(CheckpointModel::new(
            "DQN".to_string(),
            format!("checkpoint_v{}.safetensors", i + 2),
            create_dqn_epoch_50_predictor(),
        ));

        manager.stage_checkpoint("DQN", checkpoint).await.unwrap();

        let swap_latency = manager.commit_swap("DQN").await.unwrap();
        swap_latencies.push(swap_latency.as_micros() as u64);
    }

    // Calculate statistics
    swap_latencies.sort_unstable();
    let avg_latency = swap_latencies.iter().sum::<u64>() / swap_latencies.len() as u64;
    let p50_latency = swap_latencies[50];
    let p99_latency = swap_latencies[99];

    println!("Swap latency statistics (100 swaps):");
    println!("  - Average: {}μs", avg_latency);
    println!("  - P50: {}μs", p50_latency);
    println!("  - P99: {}μs", p99_latency);
    println!("  - Min: {}μs", swap_latencies[0]);
    println!("  - Max: {}μs", swap_latencies[99]);

    // All swaps should be < 1μs (but we allow 100μs for CI/testing)
    assert!(p99_latency < 100, "P99 swap latency {}μs exceeds 100μs", p99_latency);

    println!("\n=== Swap Latency Benchmark PASSED ===\n");
}

#[tokio::test]
async fn test_zero_dropped_predictions() {
    println!("\n=== Zero Dropped Predictions Test ===\n");

    let validator = CheckpointValidator::new();
    let policy = RollbackPolicy::default();
    let manager = HotSwapManager::new(validator, policy);

    // Register initial checkpoint
    let checkpoint_v30 = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "checkpoint_epoch_30.safetensors".to_string(),
        create_dqn_epoch_30_predictor(),
    ));

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

    // Spawn prediction workload (1000 predictions in background)
    let manager_clone = Arc::new(manager);
    let prediction_task = {
        let manager = manager_clone.clone();
        tokio::spawn(async move {
            let mut success_count = 0;
            let mut error_count = 0;

            for i in 0..1000 {
                let features = Features::new(
                    vec![
                        (i as f64 * 0.01).sin(),
                        (i as f64 * 0.02).cos(),
                        (i as f64 * 0.03).tanh(),
                        (i as f64 * 0.01 + 1.0).ln(),
                        (i as f64 * 0.001).exp().min(10.0),
                    ],
                    vec!["f1".to_string(), "f2".to_string(), "f3".to_string(), "f4".to_string(), "f5".to_string()],
                );

                match manager.get_active_checkpoint("DQN").await {
                    Ok(checkpoint) => match checkpoint.predict(&features) {
                        Ok(_) => success_count += 1,
                        Err(_) => error_count += 1,
                    },
                    Err(_) => error_count += 1,
                }

                // Small delay to simulate realistic prediction rate
                tokio::time::sleep(tokio::time::Duration::from_micros(100)).await;
            }

            (success_count, error_count)
        })
    };

    // Wait a bit for predictions to start
    tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;

    // Perform hot-swap while predictions are running
    println!("Performing hot-swap during active predictions...");
    let checkpoint_v50 = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "checkpoint_epoch_50.safetensors".to_string(),
        create_dqn_epoch_50_predictor(),
    ));

    manager_clone.stage_checkpoint("DQN", checkpoint_v50).await.unwrap();
    let swap_latency = manager_clone.commit_swap("DQN").await.unwrap();
    println!("✓ Hot-swap completed in {}μs during active predictions", swap_latency.as_micros());

    // Wait for prediction task to complete
    let (success_count, error_count) = prediction_task.await.unwrap();

    println!("\nPrediction results during hot-swap:");
    println!("  - Successful: {}", success_count);
    println!("  - Errors: {}", error_count);
    println!("  - Total: {}", success_count + error_count);

    // Verify zero dropped predictions (all predictions should succeed)
    assert_eq!(error_count, 0, "Found {} dropped predictions during hot-swap", error_count);
    assert_eq!(success_count, 1000, "Expected 1000 successful predictions, got {}", success_count);

    println!("\n=== Zero Dropped Predictions Test PASSED ===\n");
}