//! 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 { 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 ); }