rustytorch/crates/specialized/rtx-cfd/tests/cuda_kernel_manager_tests.rs

//! Tests for CUDA kernel manager module storage and retrieval
//!
//! Following strict TDD methodology - tests are written first,
//! then implementation is verified to pass.

#[cfg(feature = "cuda")]
mod cuda_kernel_manager_tests {
    use rtx_cfd::kernels::CudaKernelManager;
    use rtx_cfd::{CfdConfig, CfdResult};

    fn create_test_config() -> CfdConfig {
        CfdConfig {
            nx: 32,
            ny: 32,
            nz: 1,
            lx: 1.0,
            ly: 1.0,
            lz: 1.0,
            dt: 0.001,
            viscosity: 0.01,
            density: 1.0,
            device_id: 0,
            ..CfdConfig::default()
        }
    }

    #[test]
    fn test_kernel_manager_creation() -> CfdResult<()> {
        // Skip test if no GPU available
        if std::env::var("SKIP_GPU_TESTS").is_ok() {
            return Ok(());
        }

        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);

        // Verify manager is created successfully
        // Context and stream should be valid (non-null references)
        let _ = manager.context();
        let _ = manager.stream();

        Ok(())
    }

    #[test]
    fn test_module_storage_and_retrieval() -> CfdResult<()> {
        // Skip test if no GPU available
        if std::env::var("SKIP_GPU_TESTS").is_ok() {
            return Ok(());
        }

        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);

        // Test that required modules are loaded
        // Getting a module should succeed and return a valid Arc reference
        let _advection_module = manager.get_module("advection_kernels")?;

        let _diffusion_module = manager.get_module("diffusion_kernels")?;

        let _poisson_module = manager.get_module("poisson_kernels")?;

        let _matrix_module = manager.get_module("matrix_kernels")?;

        Ok(())
    }

    #[test]
    fn test_module_not_found_error() -> CfdResult<()> {
        // Skip test if no GPU available
        if std::env::var("SKIP_GPU_TESTS").is_ok() {
            return Ok(());
        }

        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);

        // Test that non-existent module returns error
        let result = manager.get_module("nonexistent_module");
        assert!(result.is_err());

        let err_msg = format!("{}", result.unwrap_err());
        assert!(err_msg.contains("not found"));

        Ok(())
    }

    #[test]
    fn test_memory_allocation() -> CfdResult<()> {
        // Skip test if no GPU available
        if std::env::var("SKIP_GPU_TESTS").is_ok() {
            return Ok(());
        }

        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);

        // Test allocating GPU memory
        let size = 1024;
        let gpu_buffer = manager.allocate_f32(size)?;
        assert_eq!(gpu_buffer.len(), size);

        Ok(())
    }

    #[test]
    fn test_memory_copy_to_device() -> CfdResult<()> {
        // Skip test if no GPU available
        if std::env::var("SKIP_GPU_TESTS").is_ok() {
            return Ok(());
        }

        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);

        // Test copying data to GPU
        let host_data: Vec<f32> = (0..100).map(|i| i as f32).collect();
        let gpu_buffer = manager.copy_to_device(&host_data)?;
        assert_eq!(gpu_buffer.len(), host_data.len());

        Ok(())
    }

    #[test]
    fn test_memory_copy_round_trip() -> CfdResult<()> {
        // Skip test if no GPU available
        if std::env::var("SKIP_GPU_TESTS").is_ok() {
            return Ok(());
        }

        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);

        // Test round-trip copy: host -> device -> host
        let original_data: Vec<f32> = (0..100).map(|i| i as f32 * 0.1).collect();
        let gpu_buffer = manager.copy_to_device(&original_data)?;
        let retrieved_data = manager.copy_from_device(&gpu_buffer)?;

        // Verify data integrity
        assert_eq!(original_data.len(), retrieved_data.len());
        for (i, (&orig, &retrieved)) in original_data.iter().zip(retrieved_data.iter()).enumerate()
        {
            assert!(
                (orig - retrieved).abs() < 1e-6,
                "Data mismatch at index {}: {} != {}",
                i,
                orig,
                retrieved
            );
        }

        Ok(())
    }

    #[test]
    fn test_stream_synchronization() -> CfdResult<()> {
        // Skip test if no GPU available
        if std::env::var("SKIP_GPU_TESTS").is_ok() {
            return Ok(());
        }

        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);

        // Test that synchronization doesn't error
        manager.synchronize()?;

        Ok(())
    }
}