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, EnsembleMetrics, HotSwapManager, RollbackPolicy,
};
use ml::{Features, MLResult, ModelPrediction};
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");
}