//! Advanced Job Tracker Tests //! //! Tests cover: //! - State machine illegal transitions (security) //! - Concurrent status updates (race conditions) //! - Progress calculation precision (floating point) //! - Database trigger validation //! - Weighted progress aggregation #![allow( dead_code, unused_variables, unused_imports, clippy::unwrap_used, clippy::expect_used, clippy::indexing_slicing, clippy::manual_range_contains )] use ml_training_service::job_spawner::{Asset, JobSpawner, ModelType}; use ml_training_service::job_tracker::{JobProgress, JobStatus, JobTracker}; use sqlx::PgPool; use std::path::PathBuf; use std::sync::Arc; use uuid::Uuid; async fn run_migrations(pool: &PgPool) { sqlx::migrate!("../../migrations") .run(pool) .await .expect("Failed to run migrations"); } async fn create_test_batch(pool: &PgPool) -> (Uuid, Vec) { let spawner = JobSpawner::new(pool.clone()); let assets = vec![Asset { symbol: "TEST".to_string(), data_file: PathBuf::from("/test/TEST.parquet"), }]; let models = vec![ModelType::DQN, ModelType::PPO, ModelType::MAMBA, ModelType::TFT]; let batch = spawner.spawn_batch(assets, models).await.unwrap(); // Get child job IDs let job_ids: Vec<(Uuid,)> = sqlx::query_as( "SELECT id FROM child_jobs WHERE batch_id = $1 ORDER BY created_at" ) .bind(batch.batch_id) .fetch_all(pool) .await .unwrap(); (batch.batch_id, job_ids.into_iter().map(|r| r.0).collect()) } // ============================================================================ // STATE MACHINE VALIDATION (ILLEGAL TRANSITIONS) // ============================================================================ #[sqlx::test] async fn test_prevent_completed_to_running_transition(pool: PgPool) { run_migrations(&pool).await; let tracker = JobTracker::new(pool.clone()); let (_, job_ids) = create_test_batch(&pool).await; let job_id = job_ids[0]; // Transition to Running then Completed tracker.update_job_status(job_id, JobStatus::Running).await.unwrap(); tracker.update_job_status(job_id, JobStatus::Completed).await.unwrap(); // Try to go back to Running (illegal) let result = tracker.update_job_status(job_id, JobStatus::Running).await; assert!(result.is_err()); let err_msg = result.unwrap_err().to_string(); assert!(err_msg.contains("Invalid state transition")); } #[sqlx::test] async fn test_prevent_failed_to_pending_transition(pool: PgPool) { run_migrations(&pool).await; let tracker = JobTracker::new(pool.clone()); let (_, job_ids) = create_test_batch(&pool).await; let job_id = job_ids[0]; // Transition to Failed tracker.update_job_status(job_id, JobStatus::Failed).await.unwrap(); // Try to go back to Pending (illegal) let result = tracker.update_job_status(job_id, JobStatus::Pending).await; assert!(result.is_err()); assert!(result.unwrap_err().to_string().contains("Invalid state transition")); } #[sqlx::test] async fn test_prevent_stopped_to_running_transition(pool: PgPool) { run_migrations(&pool).await; let tracker = JobTracker::new(pool.clone()); let (_, job_ids) = create_test_batch(&pool).await; let job_id = job_ids[0]; // Transition to Stopped tracker.update_job_status(job_id, JobStatus::Stopped).await.unwrap(); // Try to restart (illegal from terminal state) let result = tracker.update_job_status(job_id, JobStatus::Running).await; assert!(result.is_err()); } #[sqlx::test] async fn test_valid_state_transitions(pool: PgPool) { run_migrations(&pool).await; let tracker = JobTracker::new(pool.clone()); let (_, job_ids) = create_test_batch(&pool).await; // Test valid transition path: Pending → Running → Paused → Running → Completed let job_id = job_ids[0]; tracker.update_job_status(job_id, JobStatus::Running).await.unwrap(); tracker.update_job_status(job_id, JobStatus::Paused).await.unwrap(); tracker.update_job_status(job_id, JobStatus::Running).await.unwrap(); tracker.update_job_status(job_id, JobStatus::Completed).await.unwrap(); // Verify final state let status: String = sqlx::query_scalar( "SELECT status FROM child_jobs WHERE id = $1" ) .bind(job_id) .fetch_one(&pool) .await .unwrap(); assert_eq!(status, "Completed"); } #[sqlx::test] async fn test_idempotent_same_state_transition(pool: PgPool) { run_migrations(&pool).await; let tracker = JobTracker::new(pool.clone()); let (_, job_ids) = create_test_batch(&pool).await; let job_id = job_ids[0]; // Transition to Running multiple times (idempotent) tracker.update_job_status(job_id, JobStatus::Running).await.unwrap(); tracker.update_job_status(job_id, JobStatus::Running).await.unwrap(); tracker.update_job_status(job_id, JobStatus::Running).await.unwrap(); let status: String = sqlx::query_scalar( "SELECT status FROM child_jobs WHERE id = $1" ) .bind(job_id) .fetch_one(&pool) .await .unwrap(); assert_eq!(status, "Running"); } #[sqlx::test] async fn test_all_terminal_states_prevent_transitions(pool: PgPool) { run_migrations(&pool).await; let tracker = JobTracker::new(pool.clone()); let (_, job_ids) = create_test_batch(&pool).await; // Test Completed (terminal) let job1 = job_ids[0]; tracker.update_job_status(job1, JobStatus::Completed).await.unwrap(); assert!(tracker.update_job_status(job1, JobStatus::Running).await.is_err()); // Test Failed (terminal) let job2 = job_ids[1]; tracker.update_job_status(job2, JobStatus::Failed).await.unwrap(); assert!(tracker.update_job_status(job2, JobStatus::Running).await.is_err()); // Test Stopped (terminal) let job3 = job_ids[2]; tracker.update_job_status(job3, JobStatus::Stopped).await.unwrap(); assert!(tracker.update_job_status(job3, JobStatus::Running).await.is_err()); } // ============================================================================ // CONCURRENT STATUS UPDATES (RACE CONDITIONS) // ============================================================================ #[sqlx::test] async fn test_concurrent_status_updates_same_job(pool: PgPool) { run_migrations(&pool).await; let tracker = Arc::new(JobTracker::new(pool.clone())); let (_, job_ids) = create_test_batch(&pool).await; let job_id = job_ids[0]; // Transition to Running first tracker.update_job_status(job_id, JobStatus::Running).await.unwrap(); let mut handles = vec![]; // 10 concurrent attempts to complete the same job for _ in 0..10 { let tracker_clone = Arc::clone(&tracker); let handle = tokio::spawn(async move { tracker_clone .update_job_status(job_id, JobStatus::Completed) .await }); handles.push(handle); } // All should succeed (idempotent) for handle in handles { let result = handle.await.unwrap(); assert!(result.is_ok()); } // Verify final state let status: String = sqlx::query_scalar( "SELECT status FROM child_jobs WHERE id = $1" ) .bind(job_id) .fetch_one(&pool) .await .unwrap(); assert_eq!(status, "Completed"); } #[sqlx::test] async fn test_concurrent_different_status_updates(pool: PgPool) { run_migrations(&pool).await; let tracker = Arc::new(JobTracker::new(pool.clone())); let (_, job_ids) = create_test_batch(&pool).await; let job_id = job_ids[0]; // Transition to Running tracker.update_job_status(job_id, JobStatus::Running).await.unwrap(); let mut handles = vec![]; // Concurrent competing transitions for i in 0..20 { let tracker_clone = Arc::clone(&tracker); let status = if i % 3 == 0 { JobStatus::Completed } else if i % 3 == 1 { JobStatus::Failed } else { JobStatus::Paused }; let handle = tokio::spawn(async move { tracker_clone.update_job_status(job_id, status).await }); handles.push(handle); } // At least some should succeed let mut success_count = 0; for handle in handles { if handle.await.unwrap().is_ok() { success_count += 1; } } assert!(success_count > 0); // Final state should be one of the valid terminal states let status: String = sqlx::query_scalar( "SELECT status FROM child_jobs WHERE id = $1" ) .bind(job_id) .fetch_one(&pool) .await .unwrap(); assert!( status == "Completed" || status == "Failed" || status == "Paused", "Final status: {}", status ); } #[sqlx::test] async fn test_concurrent_progress_updates(pool: PgPool) { run_migrations(&pool).await; let tracker = Arc::new(JobTracker::new(pool.clone())); let (_, job_ids) = create_test_batch(&pool).await; let job_id = job_ids[0]; // Start job tracker.update_job_status(job_id, JobStatus::Running).await.unwrap(); let mut handles = vec![]; // 50 concurrent progress updates for epoch in 1..=50 { let tracker_clone = Arc::clone(&tracker); let handle = tokio::spawn(async move { tracker_clone .update_job_progress(job_id, JobProgress { job_id, current_epoch: epoch, total_epochs: 100, progress_pct: (epoch as f64 / 100.0) * 100.0, }) .await }); handles.push(handle); } // All should succeed for handle in handles { handle.await.unwrap().unwrap(); } // Final progress should be valid let (current_epoch, progress): (i32, f64) = sqlx::query_as( "SELECT current_epoch, progress_pct FROM child_jobs WHERE id = $1" ) .bind(job_id) .fetch_one(&pool) .await .unwrap(); assert!(current_epoch >= 1 && current_epoch <= 50); assert!(progress >= 0.0 && progress <= 100.0); } // ============================================================================ // PROGRESS CALCULATION PRECISION // ============================================================================ #[sqlx::test] async fn test_progress_calculation_accuracy(pool: PgPool) { run_migrations(&pool).await; let tracker = JobTracker::new(pool.clone()); let (_, job_ids) = create_test_batch(&pool).await; let job_id = job_ids[0]; // Update progress: epoch 25 of 100 tracker .update_job_progress(job_id, JobProgress { job_id, current_epoch: 25, total_epochs: 100, progress_pct: 25.0, }) .await .unwrap(); let progress: f64 = sqlx::query_scalar( "SELECT progress_pct FROM child_jobs WHERE id = $1" ) .bind(job_id) .fetch_one(&pool) .await .unwrap(); // Should be exactly 25.0 assert!((progress - 25.0).abs() < 0.01); } #[sqlx::test] async fn test_progress_boundary_conditions(pool: PgPool) { run_migrations(&pool).await; let tracker = JobTracker::new(pool.clone()); let (_, job_ids) = create_test_batch(&pool).await; // Test 0% tracker.update_job_progress(job_ids[0], JobProgress { job_id: job_ids[0], current_epoch: 0, total_epochs: 100, progress_pct: 0.0, }).await.unwrap(); let p0: f64 = sqlx::query_scalar("SELECT progress_pct FROM child_jobs WHERE id = $1") .bind(job_ids[0]) .fetch_one(&pool) .await .unwrap(); assert_eq!(p0, 0.0); // Test 100% tracker.update_job_progress(job_ids[1], JobProgress { job_id: job_ids[1], current_epoch: 100, total_epochs: 100, progress_pct: 100.0, }).await.unwrap(); let p100: f64 = sqlx::query_scalar("SELECT progress_pct FROM child_jobs WHERE id = $1") .bind(job_ids[1]) .fetch_one(&pool) .await .unwrap(); assert_eq!(p100, 100.0); } #[sqlx::test] async fn test_fractional_progress_precision(pool: PgPool) { run_migrations(&pool).await; let tracker = JobTracker::new(pool.clone()); let (_, job_ids) = create_test_batch(&pool).await; let job_id = job_ids[0]; // 1 of 3 epochs = 33.333...% tracker.update_job_progress(job_id, JobProgress { job_id, current_epoch: 1, total_epochs: 3, progress_pct: (1.0 / 3.0) * 100.0, }).await.unwrap(); let progress: f64 = sqlx::query_scalar( "SELECT progress_pct FROM child_jobs WHERE id = $1" ) .bind(job_id) .fetch_one(&pool) .await .unwrap(); // Should be close to 33.33 assert!((progress - 33.333333).abs() < 0.01); } #[sqlx::test] async fn test_weighted_batch_progress_calculation(pool: PgPool) { run_migrations(&pool).await; let tracker = JobTracker::new(pool.clone()); let (batch_id, job_ids) = create_test_batch(&pool).await; // Update progress for all 4 jobs tracker.update_job_status(job_ids[0], JobStatus::Running).await.unwrap(); tracker.update_job_progress(job_ids[0], JobProgress { job_id: job_ids[0], current_epoch: 50, total_epochs: 100, progress_pct: 50.0, }).await.unwrap(); // 50% tracker.update_job_status(job_ids[1], JobStatus::Completed).await.unwrap(); tracker.update_job_progress(job_ids[1], JobProgress { job_id: job_ids[1], current_epoch: 100, total_epochs: 100, progress_pct: 100.0, }).await.unwrap(); // 100% tracker.update_job_status(job_ids[2], JobStatus::Running).await.unwrap(); tracker.update_job_progress(job_ids[2], JobProgress { job_id: job_ids[2], current_epoch: 25, total_epochs: 100, progress_pct: 25.0, }).await.unwrap(); // 25% // Job 4 still pending (0%) // Get batch summary let summary = tracker.get_batch_summary(batch_id).await.unwrap(); // Weighted average (assuming equal weights): (50 + 100 + 25 + 0) / 4 = 43.75% assert!(summary.overall_progress >= 40.0 && summary.overall_progress <= 50.0); } // ============================================================================ // BATCH AGGREGATION VALIDATION // ============================================================================ #[sqlx::test] async fn test_batch_summary_job_counts(pool: PgPool) { run_migrations(&pool).await; let tracker = JobTracker::new(pool.clone()); let (batch_id, job_ids) = create_test_batch(&pool).await; // Update statuses: 1 running, 1 completed, 1 failed, 1 pending tracker.update_job_status(job_ids[0], JobStatus::Running).await.unwrap(); tracker.update_job_status(job_ids[1], JobStatus::Completed).await.unwrap(); tracker.update_job_status(job_ids[2], JobStatus::Failed).await.unwrap(); // job_ids[3] stays Pending let summary = tracker.get_batch_summary(batch_id).await.unwrap(); assert_eq!(summary.total_jobs, 4); assert_eq!(summary.pending_jobs, 1); assert_eq!(summary.running_jobs, 1); assert_eq!(summary.completed_jobs, 1); assert_eq!(summary.failed_jobs, 1); } #[sqlx::test] async fn test_batch_status_inference(pool: PgPool) { run_migrations(&pool).await; let tracker = JobTracker::new(pool.clone()); let (batch_id, job_ids) = create_test_batch(&pool).await; // All completed for job_id in &job_ids { tracker.update_job_status(*job_id, JobStatus::Completed).await.unwrap(); } let summary = tracker.get_batch_summary(batch_id).await.unwrap(); assert_eq!(summary.status, JobStatus::Completed); } #[sqlx::test] async fn test_batch_status_with_partial_completion(pool: PgPool) { run_migrations(&pool).await; let tracker = JobTracker::new(pool.clone()); let (batch_id, job_ids) = create_test_batch(&pool).await; // 2 completed, 2 pending tracker.update_job_status(job_ids[0], JobStatus::Completed).await.unwrap(); tracker.update_job_status(job_ids[1], JobStatus::Completed).await.unwrap(); let summary = tracker.get_batch_summary(batch_id).await.unwrap(); assert_ne!(summary.status, JobStatus::Completed); // Not all completed } // ============================================================================ // ERROR HANDLING // ============================================================================ #[sqlx::test] async fn test_update_nonexistent_job(pool: PgPool) { run_migrations(&pool).await; let tracker = JobTracker::new(pool.clone()); let fake_id = Uuid::new_v4(); let result = tracker.update_job_status(fake_id, JobStatus::Running).await; assert!(result.is_err()); } #[sqlx::test] async fn test_progress_update_without_status_running(pool: PgPool) { run_migrations(&pool).await; let tracker = JobTracker::new(pool.clone()); let (_, job_ids) = create_test_batch(&pool).await; let job_id = job_ids[0]; // Try to update progress while still Pending let result = tracker.update_job_progress(job_id, JobProgress { job_id, current_epoch: 10, total_epochs: 100, progress_pct: 10.0, }).await; // May succeed or fail depending on implementation // Just ensure no panic let _ = result; } #[sqlx::test] async fn test_invalid_progress_values(pool: PgPool) { run_migrations(&pool).await; let tracker = JobTracker::new(pool.clone()); let (_, job_ids) = create_test_batch(&pool).await; let job_id = job_ids[0]; tracker.update_job_status(job_id, JobStatus::Running).await.unwrap(); // Test epoch > total (should clamp to 100%) let result = tracker.update_job_progress(job_id, JobProgress { job_id, current_epoch: 150, total_epochs: 100, progress_pct: 150.0, }).await; assert!(result.is_ok()); let progress: f64 = sqlx::query_scalar( "SELECT progress_pct FROM child_jobs WHERE id = $1" ) .bind(job_id) .fetch_one(&pool) .await .unwrap(); // Should be clamped to 100% assert!(progress <= 100.0); } #[sqlx::test] async fn test_zero_total_epochs_handling(pool: PgPool) { run_migrations(&pool).await; let tracker = JobTracker::new(pool.clone()); let (_, job_ids) = create_test_batch(&pool).await; let job_id = job_ids[0]; tracker.update_job_status(job_id, JobStatus::Running).await.unwrap(); // Division by zero protection let result = tracker.update_job_progress(job_id, JobProgress { job_id, current_epoch: 0, total_epochs: 0, progress_pct: 0.0, }).await; // Should handle gracefully assert!(result.is_ok() || result.is_err()); // No panic }