foxhunt/services/ml_training_service/tests/gpu_resource_tests.rs

//! GPU Resource Manager Tests (TDD - Write Tests First)
//!
//! Test suite for GPU reservation system to prevent concurrent training conflicts.
//! These tests should FAIL initially, then pass after implementation.

use std::sync::Arc;
use std::time::Duration;
use tokio::time::sleep;
use uuid::Uuid;

use ml_training_service::gpu_resource_manager::{
    GPUAllocationError, GPUResourceManager,
};

/// Test 1: GPU lock acquisition should succeed when GPU is available
#[tokio::test]
async fn test_gpu_lock_acquisition_success() {
    let manager = Arc::new(GPUResourceManager::new(vec![0]).await.unwrap());
    let job_id = Uuid::new_v4();

    // Should succeed - GPU 0 is available
    let lock = manager.acquire_gpu(job_id, 0).await;
    assert!(
        lock.is_ok(),
        "GPU lock acquisition should succeed when GPU is available"
    );

    let lock = lock.unwrap();
    assert_eq!(lock.gpu_id(), 0);
    assert_eq!(lock.job_id(), job_id);
    assert!(lock.is_locked());
}

/// Test 2: GPU lock acquisition should fail when GPU is already locked
#[tokio::test]
async fn test_gpu_lock_acquisition_blocked_by_concurrent_job() {
    let manager = Arc::new(GPUResourceManager::new(vec![0]).await.unwrap());
    let job_id_1 = Uuid::new_v4();
    let job_id_2 = Uuid::new_v4();

    // First job acquires GPU 0
    let lock1 = manager.acquire_gpu(job_id_1, 0).await.unwrap();
    assert!(lock1.is_locked());

    // Second job should fail to acquire GPU 0
    let lock2 = manager.acquire_gpu(job_id_2, 0).await;
    assert!(
        lock2.is_err(),
        "Second job should fail to acquire already-locked GPU"
    );

    match lock2.unwrap_err() {
        GPUAllocationError::GPUAlreadyLocked {
            gpu_id,
            current_job_id,
        } => {
            assert_eq!(gpu_id, 0);
            assert_eq!(current_job_id, job_id_1);
        },
        _ => panic!("Expected GPUAlreadyLocked error"),
    }
}

/// Test 3: GPU lock should be released automatically on drop
#[tokio::test]
async fn test_gpu_lock_automatic_release_on_drop() {
    let manager = Arc::new(GPUResourceManager::new(vec![0]).await.unwrap());
    let job_id_1 = Uuid::new_v4();
    let job_id_2 = Uuid::new_v4();

    // First job acquires and releases GPU
    {
        let lock1 = manager.acquire_gpu(job_id_1, 0).await.unwrap();
        assert!(lock1.is_locked());
        // lock1 drops here
    }

    // Give time for cleanup
    sleep(Duration::from_millis(50)).await;

    // Second job should now succeed
    let lock2 = manager.acquire_gpu(job_id_2, 0).await;
    assert!(
        lock2.is_ok(),
        "GPU should be available after first job releases lock"
    );
}

/// Test 4: GPU memory tracking should return accurate memory usage
#[tokio::test]
async fn test_gpu_memory_tracking() {
    let manager = Arc::new(GPUResourceManager::new(vec![0]).await.unwrap());

    // Get memory info for GPU 0
    let memory_info = manager.get_gpu_memory(0).await;
    assert!(memory_info.is_ok(), "Should be able to query GPU memory");

    let memory_info = memory_info.unwrap();
    assert!(memory_info.total_mb > 0, "Total memory should be positive");
    assert!(
        memory_info.used_mb + memory_info.free_mb <= memory_info.total_mb,
        "Used + free should not exceed total"
    );
}

/// Test 5: GPU lock should be released on job crash/panic
#[tokio::test]
async fn test_gpu_lock_release_on_crash() {
    let manager = Arc::new(GPUResourceManager::new(vec![0]).await.unwrap());
    let job_id = Uuid::new_v4();

    // Simulate job crash by explicitly releasing
    let lock = manager.acquire_gpu(job_id, 0).await.unwrap();
    let gpu_id = lock.gpu_id();
    drop(lock); // Explicit drop simulates crash cleanup

    sleep(Duration::from_millis(50)).await;

    // GPU should be available again
    let new_job_id = Uuid::new_v4();
    let new_lock = manager.acquire_gpu(new_job_id, gpu_id).await;
    assert!(new_lock.is_ok(), "GPU should be released after crash");
}

/// Test 6: Multiple GPUs should support concurrent jobs
#[tokio::test]
async fn test_multiple_gpus_concurrent_jobs() {
    let manager = Arc::new(GPUResourceManager::new(vec![0, 1]).await.unwrap());
    let job_id_1 = Uuid::new_v4();
    let job_id_2 = Uuid::new_v4();

    // Two jobs on different GPUs should both succeed
    let lock1 = manager.acquire_gpu(job_id_1, 0).await;
    let lock2 = manager.acquire_gpu(job_id_2, 1).await;

    assert!(lock1.is_ok(), "First job should acquire GPU 0");
    assert!(lock2.is_ok(), "Second job should acquire GPU 1");
}

/// Test 7: Dynamic GPU allocation should select available GPU
#[tokio::test]
async fn test_dynamic_gpu_allocation() {
    let manager = Arc::new(GPUResourceManager::new(vec![0, 1]).await.unwrap());
    let job_id = Uuid::new_v4();

    // Request any available GPU (None = auto-select)
    let lock = manager.acquire_any_available_gpu(job_id).await;
    assert!(lock.is_ok(), "Should allocate an available GPU");

    let lock = lock.unwrap();
    assert!(
        lock.gpu_id() == 0 || lock.gpu_id() == 1,
        "Should allocate GPU 0 or 1"
    );
}

/// Test 8: Should reject invalid GPU IDs
#[tokio::test]
async fn test_invalid_gpu_id_rejection() {
    let manager = Arc::new(GPUResourceManager::new(vec![0]).await.unwrap());
    let job_id = Uuid::new_v4();

    // Request non-existent GPU 99
    let lock = manager.acquire_gpu(job_id, 99).await;
    assert!(lock.is_err(), "Should reject invalid GPU ID");

    match lock.unwrap_err() {
        GPUAllocationError::GPUNotFound { gpu_id } => {
            assert_eq!(gpu_id, 99);
        },
        _ => panic!("Expected GPUNotFound error"),
    }
}

/// Test 9: Explicit release should free GPU immediately
#[tokio::test]
async fn test_explicit_gpu_release() {
    let manager = Arc::new(GPUResourceManager::new(vec![0]).await.unwrap());
    let job_id = Uuid::new_v4();

    let lock = manager.acquire_gpu(job_id, 0).await.unwrap();
    let gpu_id = lock.gpu_id();

    // Explicitly release GPU
    manager.release_gpu(gpu_id, job_id).await.unwrap();

    // GPU should be immediately available
    let new_job_id = Uuid::new_v4();
    let new_lock = manager.acquire_gpu(new_job_id, gpu_id).await;
    assert!(
        new_lock.is_ok(),
        "GPU should be available after explicit release"
    );
}

/// Test 10: Concurrent acquisition attempts should be serialized
#[tokio::test]
async fn test_concurrent_acquisition_serialization() {
    let manager = Arc::new(GPUResourceManager::new(vec![0]).await.unwrap());
    let manager_clone = Arc::clone(&manager);

    // Spawn 10 concurrent tasks trying to acquire GPU 0
    let mut handles = vec![];
    for _ in 0..10 {
        let mgr = Arc::clone(&manager);
        let handle = tokio::spawn(async move {
            let job_id = Uuid::new_v4();
            mgr.acquire_gpu(job_id, 0).await
        });
        handles.push(handle);
    }

    // Collect results
    let mut successes = 0;
    let mut failures = 0;
    for handle in handles {
        match handle.await.unwrap() {
            Ok(_) => successes += 1,
            Err(_) => failures += 1,
        }
    }

    // Exactly 1 should succeed, 9 should fail
    assert_eq!(successes, 1, "Exactly one task should acquire the GPU");
    assert_eq!(failures, 9, "Nine tasks should fail to acquire the GPU");

    // Clean up by releasing all
    manager_clone.release_all().await.unwrap();
}

/// Test 11: Load test - 100 concurrent job attempts
#[tokio::test]
async fn test_load_100_concurrent_jobs() {
    let manager = Arc::new(GPUResourceManager::new(vec![0, 1, 2, 3]).await.unwrap());

    // Spawn 100 concurrent tasks
    let mut handles = vec![];
    for _ in 0..100 {
        let mgr = Arc::clone(&manager);
        let handle = tokio::spawn(async move {
            let job_id = Uuid::new_v4();
            mgr.acquire_any_available_gpu(job_id).await
        });
        handles.push(handle);
    }

    // Collect results
    let mut successes = 0;
    for handle in handles {
        if handle.await.unwrap().is_ok() {
            successes += 1;
        }
    }

    // At most 4 should succeed (4 GPUs available)
    assert!(
        successes <= 4,
        "At most 4 jobs should acquire GPUs (4 available)"
    );
    assert!(successes > 0, "At least one job should succeed");
}

/// Test 12: List active jobs on GPU
#[tokio::test]
async fn test_list_active_jobs() {
    let manager = Arc::new(GPUResourceManager::new(vec![0, 1]).await.unwrap());
    let job_id_1 = Uuid::new_v4();
    let job_id_2 = Uuid::new_v4();

    // Acquire GPUs
    let _lock1 = manager.acquire_gpu(job_id_1, 0).await.unwrap();
    let _lock2 = manager.acquire_gpu(job_id_2, 1).await.unwrap();

    // List active jobs
    let active_jobs = manager.list_active_jobs().await.unwrap();
    assert_eq!(active_jobs.len(), 2, "Should have 2 active jobs");
    assert!(active_jobs.contains(&(0, job_id_1)));
    assert!(active_jobs.contains(&(1, job_id_2)));
}

/// Test 13: GPU utilization percentage tracking
#[tokio::test]
async fn test_gpu_utilization_tracking() {
    let manager = Arc::new(GPUResourceManager::new(vec![0]).await.unwrap());

    let utilization = manager.get_gpu_utilization(0).await;
    assert!(
        utilization.is_ok(),
        "Should be able to query GPU utilization"
    );

    let utilization = utilization.unwrap();
    assert!(
        utilization >= 0.0 && utilization <= 100.0,
        "Utilization should be 0-100%"
    );
}

/// Test 14: Memory threshold enforcement
#[tokio::test]
async fn test_memory_threshold_enforcement() {
    let manager = Arc::new(GPUResourceManager::new(vec![0]).await.unwrap());
    let job_id = Uuid::new_v4();

    // Try to acquire GPU with impossible memory requirement
    let lock = manager
        .acquire_gpu_with_memory_requirement(job_id, 0, 999_999_999)
        .await;

    // Should fail if memory requirement exceeds available memory
    // (This may pass if GPU has >1TB memory, but unlikely)
    if lock.is_err() {
        match lock.unwrap_err() {
            GPUAllocationError::InsufficientMemory {
                required_mb,
                available_mb,
                ..
            } => {
                assert!(required_mb > available_mb);
            },
            _ => panic!("Expected InsufficientMemory error"),
        }
    }
}

/// Test 15: Cleanup all locks
#[tokio::test]
async fn test_cleanup_all_locks() {
    let manager = Arc::new(GPUResourceManager::new(vec![0, 1]).await.unwrap());
    let job_id_1 = Uuid::new_v4();
    let job_id_2 = Uuid::new_v4();

    // Acquire both GPUs
    let _lock1 = manager.acquire_gpu(job_id_1, 0).await.unwrap();
    let _lock2 = manager.acquire_gpu(job_id_2, 1).await.unwrap();

    // Release all
    manager.release_all().await.unwrap();

    sleep(Duration::from_millis(50)).await;

    // Both GPUs should be available
    let new_job_id = Uuid::new_v4();
    let lock = manager.acquire_gpu(new_job_id, 0).await;
    assert!(lock.is_ok(), "GPU 0 should be available after release_all");

    let lock2 = manager.acquire_gpu(Uuid::new_v4(), 1).await;
    assert!(lock2.is_ok(), "GPU 1 should be available after release_all");
}