foxhunt/services/ml_training_service/tests/job_spawner_test.rs

//! Job Spawner Tests
//!
//! Comprehensive test suite for the JobSpawner module following TDD principles.
//! Tests cover batch creation, child job spawning, ordering, rollback, and edge cases.

use ml_training_service::job_spawner::{Asset, ChildJob, JobSpawner, ModelType};
use sqlx::PgPool;
use std::path::PathBuf;
use uuid::Uuid;

/// Helper function to run migrations on test database
async fn run_migrations(pool: &PgPool) {
    sqlx::migrate!("../../migrations")
        .run(pool)
        .await
        .expect("Failed to run migrations");
}

/// Helper function to create test assets
fn create_test_assets(count: usize) -> Vec<Asset> {
    (0..count)
        .map(|i| Asset {
            symbol: format!("TEST{}", i),
            data_file: PathBuf::from(format!("/test/data/TEST{}.parquet", i)),
        })
        .collect()
}

/// Helper function to create test models in correct training order
fn create_test_models() -> Vec<ModelType> {
    vec![
        ModelType::DQN,
        ModelType::PPO,
        ModelType::MAMBA,
        ModelType::TFT,
    ]
}

#[sqlx::test]
async fn test_spawn_batch_single_asset_4_models(pool: PgPool) {
    run_migrations(&pool).await;
    let spawner = JobSpawner::new(pool.clone());
    let assets = create_test_assets(1);
    let models = create_test_models();

    let batch = spawner
        .spawn_batch(assets.clone(), models.clone())
        .await
        .expect("Failed to spawn batch");

    // Verify batch was created
    assert!(batch.batch_id != Uuid::nil());
    assert_eq!(batch.status, "Pending");
    assert_eq!(batch.assets.len(), 1);
    assert_eq!(batch.models.len(), 4);

    // Query child jobs
    let child_jobs: Vec<ChildJob> = sqlx::query_as(
        "SELECT id, batch_id, model_type, status, created_at, config_json FROM child_jobs WHERE batch_id = $1 ORDER BY created_at",
    )
    .bind(batch.batch_id)
    .fetch_all(&pool)
    .await
    .expect("Failed to fetch child jobs");

    // Should create 4 child jobs (1 asset × 4 models)
    assert_eq!(child_jobs.len(), 4);

    // Verify all jobs are Pending
    for job in &child_jobs {
        assert_eq!(job.status, "Pending");
        assert_eq!(job.batch_id, batch.batch_id);
    }
}

#[sqlx::test]
async fn test_spawn_batch_2_assets_4_models(pool: PgPool) {
    run_migrations(&pool).await;
    let spawner = JobSpawner::new(pool.clone());
    let assets = create_test_assets(2);
    let models = create_test_models();

    let batch = spawner
        .spawn_batch(assets, models)
        .await
        .expect("Failed to spawn batch");

    // Query child jobs
    let child_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM child_jobs WHERE batch_id = $1")
        .bind(batch.batch_id)
        .fetch_one(&pool)
        .await
        .expect("Failed to count jobs");

    // Should create 8 child jobs (2 assets × 4 models)
    assert_eq!(child_count, 8);

    // Verify batch metadata
    let batch_record: (i32,) = sqlx::query_as("SELECT total_jobs FROM batch_jobs WHERE id = $1")
        .bind(batch.batch_id)
        .fetch_one(&pool)
        .await
        .expect("Failed to fetch batch");

    assert_eq!(batch_record.0, 8);
}

#[sqlx::test]
async fn test_spawn_batch_sequential_order(pool: PgPool) {
    run_migrations(&pool).await;
    let spawner = JobSpawner::new(pool.clone());
    let assets = create_test_assets(1);
    let models = create_test_models();

    let batch = spawner
        .spawn_batch(assets, models)
        .await
        .expect("Failed to spawn batch");

    // Fetch jobs in creation order
    let jobs: Vec<(String,)> = sqlx::query_as(
        "SELECT model_type FROM child_jobs WHERE batch_id = $1 ORDER BY created_at",
    )
    .bind(batch.batch_id)
    .fetch_all(&pool)
    .await
    .expect("Failed to fetch jobs");

    // Verify correct sequential order: DQN → PPO → MAMBA-2 → TFT
    assert_eq!(jobs.len(), 4);
    assert_eq!(jobs[0].0, "DQN");
    assert_eq!(jobs[1].0, "PPO");
    assert_eq!(jobs[2].0, "MAMBA-2");
    assert_eq!(jobs[3].0, "TFT");
}

#[sqlx::test]
async fn test_spawn_batch_creates_parent(pool: PgPool) {
    run_migrations(&pool).await;
    let spawner = JobSpawner::new(pool.clone());
    let assets = create_test_assets(1);
    let models = create_test_models();

    let batch = spawner
        .spawn_batch(assets.clone(), models.clone())
        .await
        .expect("Failed to spawn batch");

    // Verify batch exists in database
    let batch_record: (Uuid, String, i32) =
        sqlx::query_as("SELECT id, status, total_jobs FROM batch_jobs WHERE id = $1")
            .bind(batch.batch_id)
            .fetch_one(&pool)
            .await
            .expect("Failed to fetch batch");

    assert_eq!(batch_record.0, batch.batch_id);
    assert_eq!(batch_record.1, "Pending");
    assert_eq!(batch_record.2, 4); // 1 asset × 4 models
}

#[sqlx::test]
async fn test_spawn_batch_all_pending_status(pool: PgPool) {
    run_migrations(&pool).await;
    let spawner = JobSpawner::new(pool.clone());
    let assets = create_test_assets(2);
    let models = create_test_models();

    let batch = spawner
        .spawn_batch(assets, models)
        .await
        .expect("Failed to spawn batch");

    // Count pending jobs
    let pending_count: i64 = sqlx::query_scalar(
        "SELECT COUNT(*) FROM child_jobs WHERE batch_id = $1 AND status = 'Pending'",
    )
    .bind(batch.batch_id)
    .fetch_one(&pool)
    .await
    .expect("Failed to count pending jobs");

    // All 8 jobs should be Pending
    assert_eq!(pending_count, 8);

    // Verify batch status
    let batch_status: (String, i32, i32) = sqlx::query_as(
        "SELECT status, pending_jobs, total_jobs FROM batch_jobs WHERE id = $1",
    )
    .bind(batch.batch_id)
    .fetch_one(&pool)
    .await
    .expect("Failed to fetch batch status");

    assert_eq!(batch_status.0, "Pending");
    assert_eq!(batch_status.1, 8); // pending_jobs
    assert_eq!(batch_status.2, 8); // total_jobs
}

#[sqlx::test]
async fn test_spawn_batch_unique_job_ids(pool: PgPool) {
    run_migrations(&pool).await;
    let spawner = JobSpawner::new(pool.clone());
    let assets = create_test_assets(3);
    let models = create_test_models();

    let batch = spawner
        .spawn_batch(assets, models)
        .await
        .expect("Failed to spawn batch");

    // Fetch all job IDs
    let job_ids: Vec<(Uuid,)> =
        sqlx::query_as("SELECT id FROM child_jobs WHERE batch_id = $1")
            .bind(batch.batch_id)
            .fetch_all(&pool)
            .await
            .expect("Failed to fetch job IDs");

    // Should have 12 jobs (3 assets × 4 models)
    assert_eq!(job_ids.len(), 12);

    // Verify all IDs are unique
    let unique_ids: std::collections::HashSet<_> = job_ids.iter().map(|r| r.0).collect();
    assert_eq!(unique_ids.len(), 12);

    // Verify batch ID is also unique and not nil
    assert!(batch.batch_id != Uuid::nil());
}

#[sqlx::test]
async fn test_get_next_pending_job_returns_oldest(pool: PgPool) {
    run_migrations(&pool).await;
    let spawner = JobSpawner::new(pool.clone());
    
    // Create first batch
    let assets1 = vec![Asset {
        symbol: "ES.FUT".to_string(),
        data_file: PathBuf::from("/test/ES.parquet"),
    }];
    let models1 = vec![ModelType::DQN];
    
    let batch1 = spawner
        .spawn_batch(assets1, models1)
        .await
        .expect("Failed to spawn first batch");

    // Wait a moment to ensure different timestamps
    tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;

    // Create second batch
    let assets2 = vec![Asset {
        symbol: "NQ.FUT".to_string(),
        data_file: PathBuf::from("/test/NQ.parquet"),
    }];
    let models2 = vec![ModelType::PPO];
    
    let _batch2 = spawner
        .spawn_batch(assets2, models2)
        .await
        .expect("Failed to spawn second batch");

    // Get next pending job
    let next_job = spawner
        .get_next_pending_job()
        .await
        .expect("Failed to get next job")
        .expect("No pending job found");

    // Should return job from first batch (FIFO order)
    assert_eq!(next_job.batch_id, batch1.batch_id);
    assert_eq!(next_job.model_type, "DQN");
}

#[sqlx::test]
async fn test_get_next_pending_job_empty_queue(pool: PgPool) {
    run_migrations(&pool).await;
    let spawner = JobSpawner::new(pool.clone());

    // No jobs in queue
    let next_job = spawner
        .get_next_pending_job()
        .await
        .expect("Failed to query jobs");

    assert!(next_job.is_none());
}

#[sqlx::test]
async fn test_spawn_batch_rollback_on_error(pool: PgPool) {
    run_migrations(&pool).await;
    let spawner = JobSpawner::new(pool.clone());

    // Try to create a batch with invalid model type (this should fail)
    // We'll simulate this by creating a batch and then trying to insert
    // a child job with a NULL model_type which violates NOT NULL constraint
    
    // First, verify tables are empty
    let initial_batch_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM batch_jobs")
        .fetch_one(&pool)
        .await
        .expect("Failed to count batches");
    let initial_job_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM child_jobs")
        .fetch_one(&pool)
        .await
        .expect("Failed to count jobs");

    // Create a batch with empty assets (should fail validation)
    let empty_assets: Vec<Asset> = vec![];
    let models = create_test_models();

    let result = spawner.spawn_batch(empty_assets, models).await;

    // Should return an error
    assert!(result.is_err());

    // Verify no records were created (rollback successful)
    let final_batch_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM batch_jobs")
        .fetch_one(&pool)
        .await
        .expect("Failed to count batches");
    let final_job_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM child_jobs")
        .fetch_one(&pool)
        .await
        .expect("Failed to count jobs");

    assert_eq!(final_batch_count, initial_batch_count);
    assert_eq!(final_job_count, initial_job_count);
}

#[sqlx::test]
async fn test_spawn_batch_with_config_metadata(pool: PgPool) {
    run_migrations(&pool).await;
    let spawner = JobSpawner::new(pool.clone());
    let assets = create_test_assets(1);
    let models = create_test_models();

    let batch = spawner
        .spawn_batch(assets, models)
        .await
        .expect("Failed to spawn batch");

    // Verify config_json is populated in child jobs
    let configs: Vec<(serde_json::Value,)> =
        sqlx::query_as("SELECT config_json FROM child_jobs WHERE batch_id = $1")
            .bind(batch.batch_id)
            .fetch_all(&pool)
            .await
            .expect("Failed to fetch configs");

    // All jobs should have valid JSON config
    for (config,) in configs {
        assert!(config.is_object());
    }
}

#[sqlx::test]
async fn test_spawn_multiple_batches_independently(pool: PgPool) {
    run_migrations(&pool).await;
    let spawner = JobSpawner::new(pool.clone());

    // Create batch 1
    let assets1 = create_test_assets(1);
    let models1 = vec![ModelType::DQN, ModelType::PPO];
    let batch1 = spawner
        .spawn_batch(assets1, models1)
        .await
        .expect("Failed to spawn batch 1");

    // Create batch 2
    let assets2 = create_test_assets(2);
    let models2 = vec![ModelType::MAMBA];
    let batch2 = spawner
        .spawn_batch(assets2, models2)
        .await
        .expect("Failed to spawn batch 2");

    // Verify batches are independent
    assert_ne!(batch1.batch_id, batch2.batch_id);

    // Count jobs for each batch
    let batch1_jobs: i64 =
        sqlx::query_scalar("SELECT COUNT(*) FROM child_jobs WHERE batch_id = $1")
            .bind(batch1.batch_id)
            .fetch_one(&pool)
            .await
            .expect("Failed to count batch1 jobs");

    let batch2_jobs: i64 =
        sqlx::query_scalar("SELECT COUNT(*) FROM child_jobs WHERE batch_id = $1")
            .bind(batch2.batch_id)
            .fetch_one(&pool)
            .await
            .expect("Failed to count batch2 jobs");

    assert_eq!(batch1_jobs, 2); // 1 asset × 2 models
    assert_eq!(batch2_jobs, 2); // 2 assets × 1 model
}

#[sqlx::test]
async fn test_get_batch_status(pool: PgPool) {
    run_migrations(&pool).await;
    let spawner = JobSpawner::new(pool.clone());
    let assets = create_test_assets(1);
    let models = create_test_models();

    let batch = spawner
        .spawn_batch(assets, models)
        .await
        .expect("Failed to spawn batch");

    // Get batch status
    let status = spawner
        .get_batch_status(batch.batch_id)
        .await
        .expect("Failed to get batch status")
        .expect("Batch not found");

    assert_eq!(status.batch_id, batch.batch_id);
    assert_eq!(status.status, "Pending");
    assert_eq!(status.total_jobs, 4);
    assert_eq!(status.pending_jobs, 4);
    assert_eq!(status.completed_jobs, 0);
    assert_eq!(status.failed_jobs, 0);
}