foxhunt/services/ml_training_service/tests/job_queue_tests.rs

//! TDD Job Queue Integration Tests
//!
//! This test suite validates the job queue implementation with comprehensive
//! coverage of priority handling, GPU resource management, job cancellation,
//! and Redis persistence for crash recovery.

use ml::training_pipeline::ProductionTrainingConfig;
use ml_training_service::job_queue::{JobPriority, JobQueue};
use std::collections::HashMap;
use std::time::Duration;
use uuid::Uuid;

/// Test helper to create a test config
fn create_test_config() -> ProductionTrainingConfig {
    ProductionTrainingConfig::default()
}

/// Test helper to create tags
fn create_tags(key: &str, value: &str) -> HashMap<String, String> {
    let mut tags = HashMap::new();
    tags.insert(key.to_string(), value.to_string());
    tags
}

#[tokio::test]
async fn test_job_queue_enqueue_basic() {
    // Test: Basic job enqueue operation
    let queue = JobQueue::new(10, 1).await.expect("Failed to create queue");

    let job_id = Uuid::new_v4();
    let result = queue
        .enqueue(
            job_id,
            "DQN".to_string(),
            create_test_config(),
            "Test DQN job".to_string(),
            create_tags("env", "test"),
        )
        .await;

    assert!(result.is_ok(), "Failed to enqueue job");
}

#[tokio::test]
async fn test_job_queue_priority_ordering() {
    // Test: Jobs are dequeued in priority order (DQN/PPO before MAMBA-2/TFT)
    let queue = JobQueue::new(10, 1).await.expect("Failed to create queue");

    // Enqueue jobs in non-priority order
    let mamba_id = Uuid::new_v4();
    queue
        .enqueue(
            mamba_id,
            "MAMBA_2".to_string(),
            create_test_config(),
            "MAMBA-2 job".to_string(),
            HashMap::new(),
        )
        .await
        .unwrap();

    let dqn_id = Uuid::new_v4();
    queue
        .enqueue(
            dqn_id,
            "DQN".to_string(),
            create_test_config(),
            "DQN job".to_string(),
            HashMap::new(),
        )
        .await
        .unwrap();

    let tft_id = Uuid::new_v4();
    queue
        .enqueue(
            tft_id,
            "TFT".to_string(),
            create_test_config(),
            "TFT job".to_string(),
            HashMap::new(),
        )
        .await
        .unwrap();

    let ppo_id = Uuid::new_v4();
    queue
        .enqueue(
            ppo_id,
            "PPO".to_string(),
            create_test_config(),
            "PPO job".to_string(),
            HashMap::new(),
        )
        .await
        .unwrap();

    // Dequeue and verify priority order: DQN, PPO, then MAMBA-2, TFT
    let first = queue.dequeue().await.unwrap().expect("Expected DQN job");
    assert_eq!(
        first.job_id, dqn_id,
        "Expected DQN job first (High priority)"
    );

    let second = queue.dequeue().await.unwrap().expect("Expected PPO job");
    assert_eq!(
        second.job_id, ppo_id,
        "Expected PPO job second (High priority)"
    );

    let third = queue
        .dequeue()
        .await
        .unwrap()
        .expect("Expected MAMBA-2 job");
    assert_eq!(
        third.job_id, mamba_id,
        "Expected MAMBA-2 job third (Medium priority)"
    );

    let fourth = queue.dequeue().await.unwrap().expect("Expected TFT job");
    assert_eq!(
        fourth.job_id, tft_id,
        "Expected TFT job fourth (Medium priority)"
    );
}

#[tokio::test]
async fn test_job_queue_gpu_semaphore_single_job() {
    // Test: GPU semaphore allows only 1 job at a time
    let queue = JobQueue::new(10, 1).await.expect("Failed to create queue");

    // Enqueue two GPU jobs
    let job1_id = Uuid::new_v4();
    queue
        .enqueue(
            job1_id,
            "DQN".to_string(),
            create_test_config(),
            "GPU job 1".to_string(),
            HashMap::new(),
        )
        .await
        .unwrap();

    let job2_id = Uuid::new_v4();
    queue
        .enqueue(
            job2_id,
            "PPO".to_string(),
            create_test_config(),
            "GPU job 2".to_string(),
            HashMap::new(),
        )
        .await
        .unwrap();

    // Acquire GPU permit for first job
    let permit1 = queue
        .acquire_gpu_permit()
        .await
        .expect("Failed to acquire GPU permit");

    // Try to acquire another permit - should timeout since only 1 GPU available
    let timeout_result =
        tokio::time::timeout(Duration::from_millis(100), queue.acquire_gpu_permit()).await;

    assert!(timeout_result.is_err(), "Should timeout when GPU is busy");

    // Release permit
    drop(permit1);

    // Now second permit should succeed
    let permit2 = queue
        .acquire_gpu_permit()
        .await
        .expect("Failed to acquire second GPU permit");
    drop(permit2);
}

#[tokio::test]
async fn test_job_queue_cancellation_removes_from_queue() {
    // Test: Cancel a pending job and verify it's removed from queue
    let queue = JobQueue::new(10, 1).await.expect("Failed to create queue");

    let job_id = Uuid::new_v4();
    queue
        .enqueue(
            job_id,
            "MAMBA_2".to_string(),
            create_test_config(),
            "Cancel test".to_string(),
            HashMap::new(),
        )
        .await
        .unwrap();

    // Cancel the job
    let cancelled = queue
        .cancel_job(job_id)
        .await
        .expect("Failed to cancel job");
    assert!(cancelled, "Job should be cancelled");

    // Try to dequeue - should return None since job was cancelled
    let dequeued = queue.dequeue().await.unwrap();
    assert!(
        dequeued.is_none(),
        "Queue should be empty after cancellation"
    );
}

#[tokio::test]
async fn test_job_queue_cancellation_does_not_exist() {
    // Test: Cancelling non-existent job returns false
    let queue = JobQueue::new(10, 1).await.expect("Failed to create queue");

    let fake_job_id = Uuid::new_v4();
    let cancelled = queue
        .cancel_job(fake_job_id)
        .await
        .expect("Cancel operation failed");

    assert!(!cancelled, "Should return false for non-existent job");
}

#[tokio::test]
async fn test_job_queue_empty_dequeue() {
    // Test: Dequeue from empty queue returns None immediately
    let queue = JobQueue::new(10, 1).await.expect("Failed to create queue");

    let result = queue.dequeue().await.expect("Dequeue should not fail");

    assert!(result.is_none(), "Empty queue dequeue should return None");
}

#[tokio::test]
async fn test_job_queue_capacity_full() {
    // Test: Queue respects capacity limit
    let queue = JobQueue::new(2, 1).await.expect("Failed to create queue");

    // Fill queue to capacity
    let job1_id = Uuid::new_v4();
    queue
        .enqueue(
            job1_id,
            "DQN".to_string(),
            create_test_config(),
            "Job 1".to_string(),
            HashMap::new(),
        )
        .await
        .unwrap();

    let job2_id = Uuid::new_v4();
    queue
        .enqueue(
            job2_id,
            "PPO".to_string(),
            create_test_config(),
            "Job 2".to_string(),
            HashMap::new(),
        )
        .await
        .unwrap();

    // Try to enqueue beyond capacity - should fail immediately
    let job3_id = Uuid::new_v4();
    let result = queue
        .enqueue(
            job3_id,
            "MAMBA_2".to_string(),
            create_test_config(),
            "Job 3".to_string(),
            HashMap::new(),
        )
        .await;

    assert!(
        result.is_err(),
        "Enqueue on full queue should fail immediately"
    );
    assert!(
        result.unwrap_err().to_string().contains("capacity"),
        "Error should mention capacity"
    );
}

#[tokio::test]
async fn test_job_priority_determination() {
    // Test: Priority is correctly determined from model type
    assert_eq!(
        JobPriority::from_model_type("DQN"),
        JobPriority::High,
        "DQN should be High priority"
    );

    assert_eq!(
        JobPriority::from_model_type("PPO"),
        JobPriority::High,
        "PPO should be High priority"
    );

    assert_eq!(
        JobPriority::from_model_type("MAMBA_2"),
        JobPriority::Medium,
        "MAMBA_2 should be Medium priority"
    );

    assert_eq!(
        JobPriority::from_model_type("TFT"),
        JobPriority::Medium,
        "TFT should be Medium priority"
    );

    assert_eq!(
        JobPriority::from_model_type("TLOB"),
        JobPriority::Low,
        "TLOB should be Low priority"
    );

    assert_eq!(
        JobPriority::from_model_type("UNKNOWN_MODEL"),
        JobPriority::Low,
        "Unknown models should default to Low priority"
    );
}

#[tokio::test]
async fn test_job_queue_redis_persistence_save_load() {
    // Test: Job queue state persists to Redis and can be recovered
    let redis_url =
        std::env::var("REDIS_URL").unwrap_or_else(|_| "redis://localhost:6379".to_string());

    // Create queue and enqueue jobs
    let queue = JobQueue::with_redis(10, 1, &redis_url)
        .await
        .expect("Failed to create queue with Redis");

    let job1_id = Uuid::new_v4();
    queue
        .enqueue(
            job1_id,
            "DQN".to_string(),
            create_test_config(),
            "Persistence test job 1".to_string(),
            create_tags("persistence", "test"),
        )
        .await
        .unwrap();

    let job2_id = Uuid::new_v4();
    queue
        .enqueue(
            job2_id,
            "MAMBA_2".to_string(),
            create_test_config(),
            "Persistence test job 2".to_string(),
            create_tags("persistence", "test"),
        )
        .await
        .unwrap();

    // Manually trigger persistence
    queue
        .persist_to_redis()
        .await
        .expect("Failed to persist to Redis");

    // Create new queue instance and restore from Redis
    let recovered_queue = JobQueue::with_redis(10, 1, &redis_url)
        .await
        .expect("Failed to create recovery queue");
    recovered_queue
        .restore_from_redis()
        .await
        .expect("Failed to restore from Redis");

    // Verify jobs were recovered in correct order
    let recovered1 = recovered_queue
        .dequeue()
        .await
        .unwrap()
        .expect("Expected first recovered job");
    assert_eq!(
        recovered1.job_id, job1_id,
        "First job should be DQN (high priority)"
    );

    let recovered2 = recovered_queue
        .dequeue()
        .await
        .unwrap()
        .expect("Expected second recovered job");
    assert_eq!(recovered2.job_id, job2_id, "Second job should be MAMBA_2");
}

#[tokio::test]
async fn test_job_queue_redis_persistence_crash_recovery() {
    // Test: Simulate service crash and recovery from Redis
    let redis_url =
        std::env::var("REDIS_URL").unwrap_or_else(|_| "redis://localhost:6379".to_string());

    let test_namespace = format!("crash_test_{}", Uuid::new_v4());

    // Simulate running service
    {
        let queue = JobQueue::with_redis_namespace(10, 1, &redis_url, &test_namespace)
            .await
            .expect("Failed to create queue");

        // Enqueue some jobs
        for i in 0..3 {
            let job_id = Uuid::new_v4();
            queue
                .enqueue(
                    job_id,
                    if i % 2 == 0 { "DQN" } else { "TFT" }.to_string(),
                    create_test_config(),
                    format!("Crash test job {}", i),
                    HashMap::new(),
                )
                .await
                .unwrap();
        }

        queue.persist_to_redis().await.expect("Failed to persist");

        // Simulate crash - queue goes out of scope
    }

    // Simulate service restart - create new queue and recover
    let recovered_queue = JobQueue::with_redis_namespace(10, 1, &redis_url, &test_namespace)
        .await
        .expect("Failed to create recovery queue");

    recovered_queue
        .restore_from_redis()
        .await
        .expect("Failed to restore from Redis");

    // Verify we recovered all 3 jobs
    let mut recovered_count = 0;
    for _ in 0..3 {
        if let Ok(Some(_job)) = recovered_queue.dequeue().await {
            recovered_count += 1;
        } else {
            break;
        }
    }

    assert_eq!(recovered_count, 3, "Should recover all 3 jobs after crash");
}

#[tokio::test]
async fn test_job_queue_get_status() {
    // Test: Get status of jobs in queue
    let queue = JobQueue::new(10, 1).await.expect("Failed to create queue");

    let job_id = Uuid::new_v4();
    queue
        .enqueue(
            job_id,
            "DQN".to_string(),
            create_test_config(),
            "Status test".to_string(),
            HashMap::new(),
        )
        .await
        .unwrap();

    let status = queue
        .get_job_status(job_id)
        .await
        .expect("Failed to get status");
    assert!(status.is_some(), "Job should have status");

    let status_info = status.unwrap();
    assert_eq!(status_info.job_id, job_id);
    assert_eq!(status_info.model_type, "DQN");
    assert_eq!(status_info.status, "queued");
}

#[tokio::test]
async fn test_job_queue_list_all_jobs() {
    // Test: List all jobs in queue
    let queue = JobQueue::new(10, 1).await.expect("Failed to create queue");

    // Enqueue multiple jobs
    let mut job_ids = Vec::new();
    for i in 0..5 {
        let job_id = Uuid::new_v4();
        job_ids.push(job_id);
        queue
            .enqueue(
                job_id,
                format!("MODEL_{}", i),
                create_test_config(),
                format!("Job {}", i),
                HashMap::new(),
            )
            .await
            .unwrap();
    }

    let all_jobs = queue.list_jobs().await.expect("Failed to list jobs");
    assert_eq!(all_jobs.len(), 5, "Should list all 5 jobs");

    // Verify all job IDs are present
    for job_id in job_ids {
        assert!(
            all_jobs.iter().any(|j| j.job_id == job_id),
            "Job {} should be in list",
            job_id
        );
    }
}

#[tokio::test]
async fn test_job_queue_concurrent_enqueue() {
    // Test: Multiple concurrent enqueue operations
    let queue = JobQueue::new(100, 1).await.expect("Failed to create queue");

    let mut handles = Vec::new();

    for i in 0..10 {
        let queue_clone = queue.clone();
        let handle = tokio::spawn(async move {
            let job_id = Uuid::new_v4();
            queue_clone
                .enqueue(
                    job_id,
                    format!("MODEL_{}", i),
                    create_test_config(),
                    format!("Concurrent job {}", i),
                    HashMap::new(),
                )
                .await
                .unwrap();
        });
        handles.push(handle);
    }

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

    // Verify all 10 jobs are in queue
    let all_jobs = queue.list_jobs().await.expect("Failed to list jobs");
    assert_eq!(all_jobs.len(), 10, "Should have all 10 concurrent jobs");
}

#[tokio::test]
async fn test_job_queue_metrics() {
    // Test: Queue metrics (queue length, processing count, etc.)
    let queue = JobQueue::new(10, 1).await.expect("Failed to create queue");

    // Initial metrics
    let metrics = queue.get_metrics().await.expect("Failed to get metrics");
    assert_eq!(metrics.queued_jobs, 0);
    assert_eq!(metrics.processing_jobs, 0);
    assert_eq!(metrics.available_gpu_slots, 1);

    // Enqueue jobs
    for i in 0..3 {
        let job_id = Uuid::new_v4();
        queue
            .enqueue(
                job_id,
                "DQN".to_string(),
                create_test_config(),
                format!("Metrics test job {}", i),
                HashMap::new(),
            )
            .await
            .unwrap();
    }

    // Check metrics after enqueuing
    let metrics = queue.get_metrics().await.expect("Failed to get metrics");
    assert_eq!(metrics.queued_jobs, 3);
    assert_eq!(metrics.available_gpu_slots, 1);
}

#[tokio::test]
async fn test_job_queue_redis_connection_failure_handling() {
    // Test: Graceful handling of Redis connection failures
    let invalid_redis_url = "redis://invalid-host:9999";

    // Should fail gracefully with clear error message
    let result = JobQueue::with_redis(10, 1, invalid_redis_url).await;
    assert!(result.is_err(), "Should fail with invalid Redis URL");
}

#[tokio::test]
async fn test_job_queue_priority_starvation_prevention() {
    // Test: Lower priority jobs eventually get processed (no starvation)
    let queue = JobQueue::new(10, 1).await.expect("Failed to create queue");

    // Enqueue mix of high and low priority jobs
    let low_priority_id = Uuid::new_v4();
    queue
        .enqueue(
            low_priority_id,
            "TFT".to_string(), // Medium priority
            create_test_config(),
            "Low priority job".to_string(),
            HashMap::new(),
        )
        .await
        .unwrap();

    // Enqueue high priority jobs
    for i in 0..3 {
        let job_id = Uuid::new_v4();
        queue
            .enqueue(
                job_id,
                "DQN".to_string(), // High priority
                create_test_config(),
                format!("High priority job {}", i),
                HashMap::new(),
            )
            .await
            .unwrap();
    }

    // Process all high priority jobs first
    for _ in 0..3 {
        let job = queue
            .dequeue()
            .await
            .unwrap()
            .expect("Expected high priority job");
        assert_eq!(job.model_type, "DQN");
    }

    // Now low priority job should be dequeued
    let low_job = queue
        .dequeue()
        .await
        .unwrap()
        .expect("Expected low priority job");
    assert_eq!(low_job.job_id, low_priority_id);
}

#[tokio::test]
async fn test_job_queue_load_test_100_concurrent_submissions() {
    // Load test: 100 concurrent job submissions
    let queue = JobQueue::new(200, 1).await.expect("Failed to create queue");

    let start = std::time::Instant::now();
    let mut handles = Vec::new();

    for i in 0..100 {
        let queue_clone = queue.clone();
        let handle = tokio::spawn(async move {
            let job_id = Uuid::new_v4();
            queue_clone
                .enqueue(
                    job_id,
                    if i % 2 == 0 { "DQN" } else { "MAMBA_2" }.to_string(),
                    create_test_config(),
                    format!("Load test job {}", i),
                    create_tags("load_test", "true"),
                )
                .await
                .unwrap();
        });
        handles.push(handle);
    }

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

    let duration = start.elapsed();

    // Verify all jobs were enqueued
    let all_jobs = queue.list_jobs().await.expect("Failed to list jobs");
    assert_eq!(all_jobs.len(), 100, "Should have all 100 jobs");

    // Performance assertion: 100 submissions should complete in <5 seconds
    assert!(
        duration.as_secs() < 5,
        "100 concurrent submissions took too long: {:?}",
        duration
    );

    println!(
        "✓ Load test: 100 concurrent submissions completed in {:?}",
        duration
    );
}