foxhunt/services/ml_training_service/tests/job_tracker_test.rs

//! Job Tracker Tests
//!
//! Comprehensive test suite for batch job tracking and progress aggregation.
//! Following TDD principles - tests written FIRST before implementation.

use anyhow::Result;
use sqlx::PgPool;
use uuid::Uuid;

use ml_training_service::job_tracker::{
    JobProgress, JobStatus, JobTracker,
};

/// Setup test database with migrations
async fn setup_test_db() -> Result<PgPool> {
    let database_url = std::env::var("DATABASE_URL")
        .unwrap_or_else(|_| "postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt".to_string());
    
    let pool = PgPool::connect(&database_url).await?;
    
    // Migrations already applied to database - skip to avoid conflicts
    Ok(pool)
}

/// Create a test batch job
async fn create_test_batch_job(pool: &PgPool, batch_id: Uuid, name: &str) -> Result<()> {
    sqlx::query!(
        r#"
        INSERT INTO batch_jobs (id, name, description, status)
        VALUES ($1, $2, $3, $4)
        "#,
        batch_id,
        name,
        "Test batch job",
        "Pending"
    )
    .execute(pool)
    .await?;
    
    Ok(())
}

/// Create a test child job
async fn create_test_child_job(
    pool: &PgPool,
    job_id: Uuid,
    batch_id: Uuid,
    model_type: &str,
    model_weight: f64,
    status: &str,
    progress_pct: f64,
) -> Result<()> {
    sqlx::query!(
        r#"
        INSERT INTO child_jobs (
            id, batch_id, model_type, model_weight, status, 
            progress_pct, current_epoch, total_epochs
        )
        VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
        "#,
        job_id,
        batch_id,
        model_type,
        model_weight,
        status,
        progress_pct,
        0i32,
        100i32
    )
    .execute(pool)
    .await?;
    
    Ok(())
}

/// Clean up test data
async fn cleanup_test_data(pool: &PgPool, batch_id: Uuid) -> Result<()> {
    // Child jobs will be cascade deleted
    sqlx::query!(
        "DELETE FROM batch_jobs WHERE id = $1",
        batch_id
    )
    .execute(pool)
    .await?;
    
    Ok(())
}

#[tokio::test]
async fn test_update_job_status_pending_to_running() -> Result<()> {
    let pool = setup_test_db().await?;
    let tracker = JobTracker::new(pool.clone());
    
    let batch_id = Uuid::new_v4();
    let job_id = Uuid::new_v4();
    
    // Setup
    create_test_batch_job(&pool, batch_id, "test_pending_to_running").await?;
    create_test_child_job(&pool, job_id, batch_id, "DQN", 0.10, "Pending", 0.0).await?;
    
    // Test: Update status from Pending to Running
    tracker.update_job_status(job_id, JobStatus::Running).await?;
    
    // Verify
    let status: String = sqlx::query_scalar(
        "SELECT status FROM child_jobs WHERE id = $1"
    )
    .bind(job_id)
    .fetch_one(&pool)
    .await?;
    
    assert_eq!(status, "Running");
    
    // Cleanup
    cleanup_test_data(&pool, batch_id).await?;
    
    Ok(())
}

#[tokio::test]
async fn test_update_job_status_running_to_completed() -> Result<()> {
    let pool = setup_test_db().await?;
    let tracker = JobTracker::new(pool.clone());
    
    let batch_id = Uuid::new_v4();
    let job_id = Uuid::new_v4();
    
    // Setup
    create_test_batch_job(&pool, batch_id, "test_running_to_completed").await?;
    create_test_child_job(&pool, job_id, batch_id, "PPO", 0.30, "Running", 50.0).await?;
    
    // Test: Update status from Running to Completed
    tracker.update_job_status(job_id, JobStatus::Completed).await?;
    
    // Verify
    let status: String = sqlx::query_scalar(
        "SELECT status FROM child_jobs WHERE id = $1"
    )
    .bind(job_id)
    .fetch_one(&pool)
    .await?;
    
    assert_eq!(status, "Completed");
    
    // Verify completed_at is set
    let completed_at: Option<chrono::DateTime<chrono::Utc>> = sqlx::query_scalar(
        "SELECT completed_at FROM child_jobs WHERE id = $1"
    )
    .bind(job_id)
    .fetch_one(&pool)
    .await?;
    
    assert!(completed_at.is_some(), "completed_at should be set");
    
    // Cleanup
    cleanup_test_data(&pool, batch_id).await?;
    
    Ok(())
}

#[tokio::test]
async fn test_update_job_status_invalid_transition() -> Result<()> {
    let pool = setup_test_db().await?;
    let tracker = JobTracker::new(pool.clone());
    
    let batch_id = Uuid::new_v4();
    let job_id = Uuid::new_v4();
    
    // Setup
    create_test_batch_job(&pool, batch_id, "test_invalid_transition").await?;
    create_test_child_job(&pool, job_id, batch_id, "TFT", 0.20, "Completed", 100.0).await?;
    
    // Test: Invalid transition from Completed to Running should fail
    let result = tracker.update_job_status(job_id, JobStatus::Running).await;
    
    assert!(result.is_err(), "Should reject invalid state transition");
    
    // Verify status unchanged
    let status: String = sqlx::query_scalar(
        "SELECT status FROM child_jobs WHERE id = $1"
    )
    .bind(job_id)
    .fetch_one(&pool)
    .await?;
    
    assert_eq!(status, "Completed", "Status should remain Completed");
    
    // Cleanup
    cleanup_test_data(&pool, batch_id).await?;
    
    Ok(())
}

#[tokio::test]
async fn test_update_job_progress() -> Result<()> {
    let pool = setup_test_db().await?;
    let tracker = JobTracker::new(pool.clone());
    
    let batch_id = Uuid::new_v4();
    let job_id = Uuid::new_v4();
    
    // Setup
    create_test_batch_job(&pool, batch_id, "test_update_progress").await?;
    create_test_child_job(&pool, job_id, batch_id, "MAMBA-2", 0.40, "Running", 0.0).await?;
    
    // Test: Update progress
    let progress = JobProgress {
        job_id,
        current_epoch: 50,
        total_epochs: 100,
        progress_pct: 50.0,
    };
    
    tracker.update_job_progress(job_id, progress).await?;
    
    // Verify
    let (current_epoch, total_epochs, progress_pct): (i32, i32, f64) = sqlx::query_as(
        "SELECT current_epoch, total_epochs, progress_pct FROM child_jobs WHERE id = $1"
    )
    .bind(job_id)
    .fetch_one(&pool)
    .await?;
    
    assert_eq!(current_epoch, 50);
    assert_eq!(total_epochs, 100);
    assert_eq!(progress_pct, 50.0);
    
    // Cleanup
    cleanup_test_data(&pool, batch_id).await?;
    
    Ok(())
}

#[tokio::test]
async fn test_get_batch_summary_all_pending() -> Result<()> {
    let pool = setup_test_db().await?;
    let tracker = JobTracker::new(pool.clone());
    
    let batch_id = Uuid::new_v4();
    
    // Setup: Create batch with 4 pending jobs
    create_test_batch_job(&pool, batch_id, "test_all_pending").await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "DQN", 0.10, "Pending", 0.0).await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "PPO", 0.30, "Pending", 0.0).await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "MAMBA-2", 0.40, "Pending", 0.0).await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "TFT", 0.20, "Pending", 0.0).await?;
    
    // Test: Get batch summary
    let summary = tracker.get_batch_summary(batch_id).await?;
    
    // Verify
    assert_eq!(summary.batch_id, batch_id);
    assert_eq!(summary.total_jobs, 4);
    assert_eq!(summary.pending_jobs, 4);
    assert_eq!(summary.running_jobs, 0);
    assert_eq!(summary.completed_jobs, 0);
    assert_eq!(summary.failed_jobs, 0);
    assert_eq!(summary.overall_progress, 0.0);
    assert_eq!(summary.status, JobStatus::Pending);
    
    // Cleanup
    cleanup_test_data(&pool, batch_id).await?;
    
    Ok(())
}

#[tokio::test]
async fn test_get_batch_summary_mixed_states() -> Result<()> {
    let pool = setup_test_db().await?;
    let tracker = JobTracker::new(pool.clone());
    
    let batch_id = Uuid::new_v4();
    
    // Setup: Create batch with mixed states
    create_test_batch_job(&pool, batch_id, "test_mixed_states").await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "DQN", 0.10, "Completed", 100.0).await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "PPO", 0.30, "Running", 50.0).await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "MAMBA-2", 0.40, "Running", 25.0).await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "TFT", 0.20, "Pending", 0.0).await?;
    
    // Test: Get batch summary
    let summary = tracker.get_batch_summary(batch_id).await?;
    
    // Verify
    assert_eq!(summary.batch_id, batch_id);
    assert_eq!(summary.total_jobs, 4);
    assert_eq!(summary.pending_jobs, 1);
    assert_eq!(summary.running_jobs, 2);
    assert_eq!(summary.completed_jobs, 1);
    assert_eq!(summary.failed_jobs, 0);
    assert_eq!(summary.status, JobStatus::Running);
    
    // Weighted progress: (100*0.10) + (50*0.30) + (25*0.40) + (0*0.20) = 10 + 15 + 10 + 0 = 35.0
    assert!((summary.overall_progress - 35.0).abs() < 0.01, 
            "Expected 35.0, got {}", summary.overall_progress);
    
    // Cleanup
    cleanup_test_data(&pool, batch_id).await?;
    
    Ok(())
}

#[tokio::test]
async fn test_calculate_weighted_progress_all_complete() -> Result<()> {
    let pool = setup_test_db().await?;
    let tracker = JobTracker::new(pool.clone());
    
    let batch_id = Uuid::new_v4();
    
    // Setup: All jobs completed
    create_test_batch_job(&pool, batch_id, "test_all_complete").await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "DQN", 0.10, "Completed", 100.0).await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "PPO", 0.30, "Completed", 100.0).await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "MAMBA-2", 0.40, "Completed", 100.0).await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "TFT", 0.20, "Completed", 100.0).await?;
    
    // Test: Get batch summary
    let summary = tracker.get_batch_summary(batch_id).await?;
    
    // Verify: All complete = 100% progress
    assert_eq!(summary.overall_progress, 100.0);
    assert_eq!(summary.status, JobStatus::Completed);
    assert_eq!(summary.completed_jobs, 4);
    
    // Cleanup
    cleanup_test_data(&pool, batch_id).await?;
    
    Ok(())
}

#[tokio::test]
async fn test_calculate_weighted_progress_half_complete() -> Result<()> {
    let pool = setup_test_db().await?;
    let tracker = JobTracker::new(pool.clone());
    
    let batch_id = Uuid::new_v4();
    
    // Setup: Each job at 50% progress
    create_test_batch_job(&pool, batch_id, "test_half_complete").await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "DQN", 0.10, "Running", 50.0).await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "PPO", 0.30, "Running", 50.0).await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "MAMBA-2", 0.40, "Running", 50.0).await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "TFT", 0.20, "Running", 50.0).await?;
    
    // Test: Get batch summary
    let summary = tracker.get_batch_summary(batch_id).await?;
    
    // Verify: Weighted average = 50% (since all weights sum to 1.0)
    assert_eq!(summary.overall_progress, 50.0);
    assert_eq!(summary.status, JobStatus::Running);
    
    // Cleanup
    cleanup_test_data(&pool, batch_id).await?;
    
    Ok(())
}

#[tokio::test]
async fn test_calculate_weighted_progress_dqn_only() -> Result<()> {
    let pool = setup_test_db().await?;
    let tracker = JobTracker::new(pool.clone());
    
    let batch_id = Uuid::new_v4();
    
    // Setup: Only DQN completed (10% weight)
    create_test_batch_job(&pool, batch_id, "test_dqn_only").await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "DQN", 0.10, "Completed", 100.0).await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "PPO", 0.30, "Pending", 0.0).await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "MAMBA-2", 0.40, "Pending", 0.0).await?;
    create_test_child_job(&pool, Uuid::new_v4(), batch_id, "TFT", 0.20, "Pending", 0.0).await?;
    
    // Test: Get batch summary
    let summary = tracker.get_batch_summary(batch_id).await?;
    
    // Verify: Only DQN at 100% with 10% weight = 10.0% overall
    assert_eq!(summary.overall_progress, 10.0);
    assert_eq!(summary.completed_jobs, 1);
    assert_eq!(summary.pending_jobs, 3);
    
    // Cleanup
    cleanup_test_data(&pool, batch_id).await?;
    
    Ok(())
}

#[tokio::test]
async fn test_get_batch_summary_nonexistent_batch() -> Result<()> {
    let pool = setup_test_db().await?;
    let tracker = JobTracker::new(pool.clone());
    
    let nonexistent_batch_id = Uuid::new_v4();
    
    // Test: Query nonexistent batch should error
    let result = tracker.get_batch_summary(nonexistent_batch_id).await;
    
    assert!(result.is_err(), "Should return error for nonexistent batch");
    
    Ok(())
}

#[tokio::test]
async fn test_batch_status_transitions() -> Result<()> {
    let pool = setup_test_db().await?;
    let tracker = JobTracker::new(pool.clone());
    
    let batch_id = Uuid::new_v4();
    let job1_id = Uuid::new_v4();
    let job2_id = Uuid::new_v4();
    
    // Setup: 2 jobs
    create_test_batch_job(&pool, batch_id, "test_status_transitions").await?;
    create_test_child_job(&pool, job1_id, batch_id, "DQN", 0.50, "Pending", 0.0).await?;
    create_test_child_job(&pool, job2_id, batch_id, "PPO", 0.50, "Pending", 0.0).await?;
    
    // Test: Initially pending
    let summary = tracker.get_batch_summary(batch_id).await?;
    assert_eq!(summary.status, JobStatus::Pending);
    
    // Test: Start first job -> batch becomes Running
    tracker.update_job_status(job1_id, JobStatus::Running).await?;
    let summary = tracker.get_batch_summary(batch_id).await?;
    assert_eq!(summary.status, JobStatus::Running);
    
    // Test: Complete first job -> batch still Running
    tracker.update_job_status(job1_id, JobStatus::Completed).await?;
    let summary = tracker.get_batch_summary(batch_id).await?;
    assert_eq!(summary.status, JobStatus::Pending); // Second job still pending
    
    // Test: Start second job -> batch Running again
    tracker.update_job_status(job2_id, JobStatus::Running).await?;
    let summary = tracker.get_batch_summary(batch_id).await?;
    assert_eq!(summary.status, JobStatus::Running);
    
    // Test: Complete second job -> batch Completed
    tracker.update_job_status(job2_id, JobStatus::Completed).await?;
    let summary = tracker.get_batch_summary(batch_id).await?;
    assert_eq!(summary.status, JobStatus::Completed);
    
    // Cleanup
    cleanup_test_data(&pool, batch_id).await?;
    
    Ok(())
}

#[tokio::test]
async fn test_batch_status_with_failures() -> Result<()> {
    let pool = setup_test_db().await?;
    let tracker = JobTracker::new(pool.clone());
    
    let batch_id = Uuid::new_v4();
    let job1_id = Uuid::new_v4();
    let job2_id = Uuid::new_v4();
    
    // Setup: 2 jobs
    create_test_batch_job(&pool, batch_id, "test_failures").await?;
    create_test_child_job(&pool, job1_id, batch_id, "DQN", 0.50, "Running", 30.0).await?;
    create_test_child_job(&pool, job2_id, batch_id, "PPO", 0.50, "Running", 40.0).await?;
    
    // Test: Fail first job
    tracker.update_job_status(job1_id, JobStatus::Failed).await?;
    let summary = tracker.get_batch_summary(batch_id).await?;
    assert_eq!(summary.failed_jobs, 1);
    assert_eq!(summary.status, JobStatus::Running); // Still has running job
    
    // Test: Fail second job -> batch becomes Failed
    tracker.update_job_status(job2_id, JobStatus::Failed).await?;
    let summary = tracker.get_batch_summary(batch_id).await?;
    assert_eq!(summary.failed_jobs, 2);
    assert_eq!(summary.status, JobStatus::Failed);
    
    // Cleanup
    cleanup_test_data(&pool, batch_id).await?;
    
    Ok(())
}

#[tokio::test]
async fn test_progress_updates_trigger_batch_aggregation() -> Result<()> {
    let pool = setup_test_db().await?;
    let tracker = JobTracker::new(pool.clone());
    
    let batch_id = Uuid::new_v4();
    let job_id = Uuid::new_v4();
    
    // Setup
    create_test_batch_job(&pool, batch_id, "test_progress_aggregation").await?;
    create_test_child_job(&pool, job_id, batch_id, "MAMBA-2", 0.40, "Running", 0.0).await?;
    
    // Test: Update progress to 75%
    let progress = JobProgress {
        job_id,
        current_epoch: 75,
        total_epochs: 100,
        progress_pct: 75.0,
    };
    tracker.update_job_progress(job_id, progress).await?;
    
    // Verify: Batch summary reflects updated progress
    let summary = tracker.get_batch_summary(batch_id).await?;
    // MAMBA-2 at 75% with 0.40 weight = 30.0% overall
    assert!((summary.overall_progress - 30.0).abs() < 0.01,
            "Expected 30.0, got {}", summary.overall_progress);
    
    // Cleanup
    cleanup_test_data(&pool, batch_id).await?;
    
    Ok(())
}