foxhunt/services/ml_training_service/tests/advanced_job_tracker_tests.rs

//! 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

use ml_training_service::job_spawner::{Asset, JobSpawner, ModelType};
use ml_training_service::job_tracker::{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<Uuid>) {
    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::MAMBA2, 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() {
    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, epoch, 100)
                .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, 25, 100)
        .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], 0, 100).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], 100, 100).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, 1, 3).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], 50, 100).await.unwrap(); // 50%

    tracker.update_job_status(job_ids[1], JobStatus::Completed).await.unwrap();
    tracker.update_job_progress(job_ids[1], 100, 100).await.unwrap(); // 100%

    tracker.update_job_status(job_ids[2], JobStatus::Running).await.unwrap();
    tracker.update_job_progress(job_ids[2], 25, 100).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, 10, 100).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, 150, 100).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, 0, 0).await;

    // Should handle gracefully
    assert!(result.is_ok() || result.is_err()); // No panic
}