foxhunt/services/ml_training_service/tests/integration/failure_recovery_test.rs

//! Failure Recovery Integration Tests
//!
//! Tests error handling, reconnection, and partial batch completion
//!
//! **Test Coverage:**
//! - Job failure → Batch status update
//! - Database connection loss → Reconnect
//! - File not found → Graceful error
//! - GPU OOM → Fallback to CPU
//! - Partial batch completion

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

use ml_training_service::{
    database::DatabaseManager,
    orchestrator::{JobStatus, TrainingOrchestrator},
    storage::{ModelStorageManager, StorageConfig},
};
use tempfile::TempDir;
use uuid::Uuid;

/// Create test orchestrator
async fn setup_test_orchestrator() -> (Arc<TrainingOrchestrator>, TempDir) {
    let ml_config = config::MLConfig::default();

    let database_url = std::env::var("DATABASE_URL").unwrap_or_else(|_| {
        "postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt".to_string()
    });

    let db_config = config::database::DatabaseConfig {
        url: database_url.clone(),
        max_connections: 5,
        min_connections: 1,
        connect_timeout: Duration::from_secs(10),
        query_timeout: Duration::from_secs(30),
        enable_query_logging: false,
        application_name: Some("failure_recovery_test".to_string()),
        pool: config::PoolConfig {
            min_connections: 1,
            max_connections: 5,
            acquire_timeout_secs: 5,
            max_lifetime_secs: 3600,
            idle_timeout_secs: 300,
            test_before_acquire: true,
            database_url: database_url.clone(),
            health_check_enabled: true,
            health_check_interval_secs: 60,
        },
        transaction: config::TransactionConfig::default(),
    };

    let database = Arc::new(
        DatabaseManager::new(&db_config)
            .await
            .expect("Failed to create database"),
    );

    let temp_dir = TempDir::new().expect("Failed to create temp dir");
    let storage_config = StorageConfig {
        storage_type: "local".to_string(),
        local_base_path: Some(temp_dir.path().to_path_buf()),
        enable_compression: false,
    };
    let storage = Arc::new(
        ModelStorageManager::new(storage_config)
            .await
            .expect("Failed to create storage"),
    );

    let orchestrator = Arc::new(
        TrainingOrchestrator::new(ml_config, database, storage)
            .await
            .expect("Failed to create orchestrator"),
    );

    (orchestrator, temp_dir)
}

/// Create test config
fn create_test_config() -> ml::training_pipeline::ProductionTrainingConfig {
    use ml::training_pipeline::{
        FinancialValidationConfig, ModelArchitectureConfig, PerformanceConfig,
        ProductionTrainingConfig, TrainingHyperparameters,
    };

    ProductionTrainingConfig {
        model_config: ModelArchitectureConfig {
            input_dim: 225,
            output_dim: 1,
            hidden_dims: vec![64, 32],
            dropout_rate: 0.1,
            activation: "relu".to_string(),
            batch_norm: false,
            residual_connections: false,
        },
        training_params: TrainingHyperparameters {
            learning_rate: 0.001,
            batch_size: 32,
            max_epochs: 3,
            patience: 2,
            validation_split: 0.2,
            l2_regularization: 0.0001,
            lr_decay_factor: 0.1,
            lr_decay_patience: 2,
        },
        safety_config: ml::safety::MLSafetyConfig::default(),
        gradient_config: ml::safety::GradientSafetyConfig::default(),
        financial_config: FinancialValidationConfig {
            max_prediction_multiple: 2.0,
            min_prediction_confidence: 0.6,
            validate_position_sizing: true,
            max_position_fraction: 0.25,
            min_sharpe_threshold: 0.5,
        },
        performance_config: PerformanceConfig {
            device_preference: "cpu".to_string(),
            num_workers: 2,
            max_memory_bytes: 2 * 1024 * 1024 * 1024,
            gradient_accumulation_steps: 1,
        },
    }
}

// ============================================================================
// Test 1: Job failure → Batch status update
// ============================================================================

#[tokio::test]
async fn test_job_failure_updates_batch() {
    let (orchestrator, _temp_dir) = setup_test_orchestrator().await;

    // Submit job
    let config = create_test_config();
    let mut tags = std::collections::HashMap::new();
    tags.insert("failure_test".to_string(), "batch_update".to_string());

    let job_id = orchestrator
        .submit_job(
            "DQN".to_string(),
            config,
            "Failure test job".to_string(),
            tags,
        )
        .await
        .expect("Failed to submit job");

    println!("✅ Submitted job: {}", job_id);

    // Simulate failure by stopping with error
    let stopped = orchestrator
        .stop_job(job_id, "Simulated failure".to_string())
        .await
        .expect("Failed to stop job");

    assert!(stopped);
    println!("✅ Job stopped (simulating failure)");

    // Verify job is in Stopped state
    let job = orchestrator.get_job(job_id).await.expect("Failed to get job");
    assert_eq!(job.status, JobStatus::Stopped);
    assert!(job.error_message.is_some());
    println!("✅ Job failure recorded: {:?}", job.error_message);
}

// ============================================================================
// Test 2: Database connection loss → Reconnect
// ============================================================================

#[tokio::test]
async fn test_database_reconnection() {
    let database_url = std::env::var("DATABASE_URL").unwrap_or_else(|_| {
        "postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt".to_string()
    });

    let db_config = config::database::DatabaseConfig {
        url: database_url.clone(),
        max_connections: 5,
        min_connections: 1,
        connect_timeout: Duration::from_secs(10),
        query_timeout: Duration::from_secs(30),
        enable_query_logging: false,
        application_name: Some("reconnect_test".to_string()),
        pool: config::PoolConfig {
            min_connections: 1,
            max_connections: 5,
            acquire_timeout_secs: 5,
            max_lifetime_secs: 3600,
            idle_timeout_secs: 300,
            test_before_acquire: true,
            database_url: database_url.clone(),
            health_check_enabled: true,
            health_check_interval_secs: 60,
        },
        transaction: config::TransactionConfig::default(),
    };

    // Create first connection
    let database1 = DatabaseManager::new(&db_config)
        .await
        .expect("Failed to create database");

    println!("✅ Initial database connection established");

    // Test query
    let result = database1.health_check().await;
    assert!(result.is_ok());
    println!("✅ Health check passed");

    // Simulate reconnection by creating new manager
    let database2 = DatabaseManager::new(&db_config)
        .await
        .expect("Failed to reconnect database");

    println!("✅ Database reconnection successful");

    // Verify new connection works
    let result2 = database2.health_check().await;
    assert!(result2.is_ok());
    println!("✅ Reconnected health check passed");
}

// ============================================================================
// Test 3: File not found → Graceful error
// ============================================================================

#[tokio::test]
async fn test_file_not_found_error() {
    let (orchestrator, _temp_dir) = setup_test_orchestrator().await;

    // Try to query non-existent job
    let fake_id = Uuid::new_v4();
    let result = orchestrator.get_job(fake_id).await;

    assert!(result.is_err());
    println!("✅ Non-existent job returns error");

    let error = result.unwrap_err();
    println!("✅ Error message: {}", error);

    // Verify error contains job ID
    let error_str = error.to_string();
    assert!(error_str.contains(&fake_id.to_string()));
    println!("✅ Error contains job ID");
}

// ============================================================================
// Test 4: GPU OOM → Fallback to CPU
// ============================================================================

#[tokio::test]
async fn test_gpu_oom_cpu_fallback() {
    // Test GPU fallback configuration
    let mut config = create_test_config();

    // Start with GPU preference
    config.performance_config.device_preference = "cuda".to_string();
    println!("✅ Initial config: device={}", config.performance_config.device_preference);

    // Simulate GPU OOM by switching to CPU
    config.performance_config.device_preference = "cpu".to_string();
    config.performance_config.max_memory_bytes = 1024 * 1024 * 1024; // 1GB

    println!("✅ Fallback config: device={}", config.performance_config.device_preference);
    println!("✅ Reduced memory: {} GB", config.performance_config.max_memory_bytes / (1024 * 1024 * 1024));

    // Verify CPU config is valid
    assert_eq!(config.performance_config.device_preference, "cpu");
    println!("✅ CPU fallback validated");
}

// ============================================================================
// Test 5: Partial batch completion
// ============================================================================

#[tokio::test]
async fn test_partial_batch_completion() {
    let (orchestrator, _temp_dir) = setup_test_orchestrator().await;

    // Submit batch of 3 jobs
    let config = create_test_config();
    let mut job_ids = Vec::new();

    for i in 0..3 {
        let mut tags = std::collections::HashMap::new();
        tags.insert("batch".to_string(), "partial_test".to_string());
        tags.insert("index".to_string(), i.to_string());

        let job_id = orchestrator
            .submit_job(
                "DQN".to_string(),
                config.clone(),
                format!("Partial batch job #{}", i),
                tags,
            )
            .await
            .expect("Failed to submit job");

        job_ids.push(job_id);
    }

    println!("✅ Submitted batch of 3 jobs");

    // Simulate partial completion: stop first 2 jobs
    for (i, job_id) in job_ids.iter().take(2).enumerate() {
        let stopped = orchestrator
            .stop_job(*job_id, format!("Completed job #{}", i))
            .await
            .expect("Failed to stop job");

        assert!(stopped);
    }

    println!("✅ Stopped 2/3 jobs (partial completion)");

    // Verify batch state
    let all_jobs = orchestrator
        .list_jobs(None, None, None, None)
        .await
        .expect("Failed to list jobs");

    let batch_jobs: Vec<_> = all_jobs
        .iter()
        .filter(|job| {
            job.tags
                .get("batch")
                .map(|v| v == "partial_test")
                .unwrap_or(false)
        })
        .collect();

    assert_eq!(batch_jobs.len(), 3);

    let stopped_count = batch_jobs
        .iter()
        .filter(|job| job.status == JobStatus::Stopped)
        .count();

    let pending_count = batch_jobs
        .iter()
        .filter(|job| job.status == JobStatus::Pending)
        .count();

    assert_eq!(stopped_count, 2);
    assert_eq!(pending_count, 1);

    println!("✅ Partial completion validated:");
    println!("   - Stopped: {}", stopped_count);
    println!("   - Pending: {}", pending_count);

    // Verify partial progress
    let progress_percentage = (stopped_count as f32 / batch_jobs.len() as f32) * 100.0;
    println!("✅ Batch progress: {:.1}%", progress_percentage);
    assert!((progress_percentage - 66.67).abs() < 0.1);
}