foxhunt/services/ml_training_service/tests/trial_executor_test.rs

//! Integration tests for trial executor
//!
//! These tests verify the trial executor's ability to manage
//! concurrent trial execution with GPU resource allocation.

use ml_training_service::trial_executor::{TrialExecutor, PoolStats};

#[tokio::test]
async fn test_trial_executor_creation() {
    // Test basic creation with default pool size
    let executor = TrialExecutor::new("http://localhost:50054".to_string())
        .await
        .expect("Failed to create trial executor");

    // Should default to 1 GPU
    assert_eq!(executor.pool_size(), 1);
    assert!(!executor.is_shutting_down());

    // Shutdown gracefully
    executor
        .shutdown()
        .await
        .expect("Failed to shutdown executor");
    assert!(executor.is_shutting_down());
}

#[tokio::test]
async fn test_trial_executor_with_explicit_pool_size() {
    // Test creation with explicit pool size
    let executor = TrialExecutor::with_pool_size("http://localhost:50054".to_string(), 2)
        .await
        .expect("Failed to create trial executor");

    assert_eq!(executor.pool_size(), 2);

    // Get initial stats
    let stats = executor.get_stats().await;
    assert_eq!(stats.pool_size, 2);
    assert_eq!(stats.active_trials, 0);
    assert_eq!(stats.total_completed, 0);
    assert_eq!(stats.total_failed, 0);
    assert_eq!(stats.worker_stats.len(), 2);

    // Verify worker IDs
    for worker in &stats.worker_stats {
        assert!(worker.worker_id < 2);
        assert!(worker.gpu_id < 2);
        assert_eq!(worker.trials_completed, 0);
        assert_eq!(worker.trials_failed, 0);
        assert!(!worker.is_busy);
    }

    executor.shutdown().await.expect("Shutdown failed");
}

#[tokio::test]
async fn test_trial_executor_stats() {
    let executor = TrialExecutor::with_pool_size("http://localhost:50054".to_string(), 1)
        .await
        .expect("Failed to create executor");

    // Get stats
    let stats = executor.get_stats().await;

    // Verify initial state
    assert_eq!(stats.pool_size, 1);
    assert_eq!(stats.active_trials, 0);
    assert_eq!(stats.total_completed, 0);
    assert_eq!(stats.total_failed, 0);

    // Verify worker stats
    assert_eq!(stats.worker_stats.len(), 1);
    let worker = &stats.worker_stats[0];
    assert_eq!(worker.worker_id, 0);
    assert_eq!(worker.gpu_id, 0);
    assert_eq!(worker.trials_completed, 0);
    assert_eq!(worker.trials_failed, 0);
    assert_eq!(worker.oom_errors, 0);
    assert!(!worker.is_busy);

    executor.shutdown().await.expect("Shutdown failed");
}

#[tokio::test]
async fn test_trial_executor_zero_pool_size_error() {
    // Test that zero pool size returns error
    let result = TrialExecutor::with_pool_size("http://localhost:50054".to_string(), 0).await;
    assert!(result.is_err());
    assert!(result
        .unwrap_err()
        .to_string()
        .contains("Pool size must be at least 1"));
}

#[tokio::test]
async fn test_trial_executor_shutdown_idempotent() {
    let executor = TrialExecutor::with_pool_size("http://localhost:50054".to_string(), 1)
        .await
        .expect("Failed to create executor");

    // First shutdown
    executor.shutdown().await.expect("First shutdown failed");
    assert!(executor.is_shutting_down());

    // Second shutdown should be no-op
    executor
        .shutdown()
        .await
        .expect("Second shutdown failed");
    assert!(executor.is_shutting_down());
}

#[test]
fn test_gpu_detection_default() {
    // Without CUDA_VISIBLE_DEVICES, should default to 1
    std::env::remove_var("CUDA_VISIBLE_DEVICES");

    let rt = tokio::runtime::Runtime::new().unwrap();
    rt.block_on(async {
        let count = ml_training_service::trial_executor::TrialExecutor::detect_gpu_count().await;
        // Should be at least 1 (default)
        assert!(count >= 1);
    });
}

#[test]
fn test_gpu_detection_with_cuda_env() {
    // Test multi-GPU detection
    std::env::set_var("CUDA_VISIBLE_DEVICES", "0,1,2");

    let rt = tokio::runtime::Runtime::new().unwrap();
    rt.block_on(async {
        let count = ml_training_service::trial_executor::TrialExecutor::detect_gpu_count().await;
        assert_eq!(count, 3);
    });

    std::env::remove_var("CUDA_VISIBLE_DEVICES");
}

#[test]
fn test_gpu_detection_with_single_gpu() {
    // Test single GPU detection
    std::env::set_var("CUDA_VISIBLE_DEVICES", "0");

    let rt = tokio::runtime::Runtime::new().unwrap();
    rt.block_on(async {
        let count = ml_training_service::trial_executor::TrialExecutor::detect_gpu_count().await;
        assert_eq!(count, 1);
    });

    std::env::remove_var("CUDA_VISIBLE_DEVICES");
}

#[test]
fn test_gpu_detection_with_empty_env() {
    // Test empty CUDA_VISIBLE_DEVICES
    std::env::set_var("CUDA_VISIBLE_DEVICES", "");

    let rt = tokio::runtime::Runtime::new().unwrap();
    rt.block_on(async {
        let count = ml_training_service::trial_executor::TrialExecutor::detect_gpu_count().await;
        assert_eq!(count, 1); // Should default to 1
    });

    std::env::remove_var("CUDA_VISIBLE_DEVICES");
}