foxhunt/AGENT_23_TEST_IMPLEMENTATION_GUIDE.md

Agent 23: Test Implementation Quick Reference

Purpose: Quick reference for implementing 21 critical ML test cases Priority Order: Phase 1 (8 tests) → Phase 2 (4 tests) → Phase 3 (7 tests) → Phase 4 (2 tests)

---

Phase 1: Production Blockers (16 hours)

Test #5: Zero Batch Size Handling

File: ml/tests/tft_zero_batch_test.rs (NEW)

#[test]
fn test_tft_forward_pass_with_zero_batch_size() -> Result<(), MLError> {
    let device = Device::Cpu;
    let config = TFTConfig::default();
    let mut tft = TemporalFusionTransformer::new(config, device)?;

    let static_features = Tensor::zeros((0, 5), DType::F32, &device)?;
    let historical_features = Tensor::zeros((0, 60, 210), DType::F32, &device)?;
    let future_features = Tensor::zeros((0, 10, 10), DType::F32, &device)?;

    let output = tft.forward(&static_features, &historical_features, &future_features)?;
    assert_eq!(output.dims(), &[0, 10, 9]);
    Ok(())
}

Test #6: Batch Size Mismatch Validation

File: ml/tests/tft_batch_mismatch_test.rs (NEW)

#[test]
fn test_tft_input_validation_detects_mismatched_batch_sizes() -> Result<(), MLError> {
    let device = Device::Cpu;
    let config = TFTConfig::default();
    let mut tft = TemporalFusionTransformer::new(config, device)?;

    // Mismatched batch sizes: 4 vs 2
    let static_features = Tensor::zeros((4, 5), DType::F32, &device)?;
    let historical_features = Tensor::zeros((2, 60, 210), DType::F32, &device)?;
    let future_features = Tensor::zeros((2, 10, 10), DType::F32, &device)?;

    let result = tft.forward(&static_features, &historical_features, &future_features);
    assert!(result.is_err());
    assert!(matches!(result.unwrap_err(), MLError::ModelError(_)));
    Ok(())
}

TODO: Extend validate_input_dimensions in ml/src/tft/mod.rs to check batch consistency.

Test #8: NaN/Inf Input Propagation

File: ml/tests/tft_nan_handling_test.rs (NEW)

#[test]
fn test_tft_forward_pass_with_nan_inf_inputs() -> Result<(), MLError> {
    let device = Device::Cpu;
    let config = TFTConfig::default();
    let mut tft = TemporalFusionTransformer::new(config, device)?;

    // Create inputs with NaN
    let mut static_data = vec![0.0f32; 5];
    static_data[2] = f32::NAN;
    let static_features = Tensor::from_vec(static_data, (1, 5), &device)?;

    let historical_features = Tensor::zeros((1, 60, 210), DType::F32, &device)?;
    let future_features = Tensor::zeros((1, 10, 10), DType::F32, &device)?;

    // Should NOT panic, NaN propagates
    let output = tft.forward(&static_features, &historical_features, &future_features)?;

    // Verify NaN present in output (makes corruption detectable)
    let output_vec = output.to_vec3::<f32>()?;
    let has_nan = output_vec.iter()
        .flat_map(|batch| batch.iter())
        .flat_map(|horizon| horizon.iter())
        .any(|&val| val.is_nan());
    assert!(has_nan, "NaN should propagate through model");
    Ok(())
}

Test #9: Corrupt Checkpoint Deserialization

File: ml/tests/tft_corrupt_checkpoint_test.rs (NEW)

#[test]
fn test_tft_deserialization_of_corrupt_checkpoint_data() -> Result<(), MLError> {
    let device = Device::Cpu;
    let config = TFTConfig::default();
    let tft = TemporalFusionTransformer::new(config, device)?;

    // Test 1: Random bytes
    let random_bytes = vec![0xDE, 0xAD, 0xBE, 0xEF; 1024];
    let result = tft.deserialize_state(&random_bytes);
    assert!(result.is_err());

    // Test 2: Truncated file (incomplete safetensors)
    let truncated = vec![0x00; 50]; // Way too small
    let result = tft.deserialize_state(&truncated);
    assert!(result.is_err());

    // Test 3: Wrong format (JSON instead of safetensors)
    let json_data = br#"{"weights": [1.0, 2.0, 3.0]}"#;
    let result = tft.deserialize_state(json_data);
    assert!(result.is_err());

    Ok(())
}

Test #11: GPU OOM Handling

File: ml/tests/tft_gpu_oom_test.rs (NEW)

#[test]
#[cfg(feature = "cuda")]
#[ignore = "Requires GPU with limited VRAM"]
fn test_gpu_oom_during_forward_pass() -> Result<(), MLError> {
    let device = Device::cuda_if_available(0)?;
    if !device.is_cuda() {
        return Ok(()); // Skip if no GPU
    }

    // Allocate large tensor to consume most VRAM
    let vram_consumer = Tensor::zeros((8000, 8000), DType::F32, &device)?;

    // Attempt large batch that should trigger OOM
    let config = TFTConfig {
        input_dim: 225,
        hidden_dim: 512, // Larger than default
        ..Default::default()
    };
    let mut tft = TemporalFusionTransformer::new(config, device.clone())?;

    let static_features = Tensor::zeros((256, 5), DType::F32, &device)?; // Large batch
    let historical_features = Tensor::zeros((256, 60, 210), DType::F32, &device)?;
    let future_features = Tensor::zeros((256, 10, 10), DType::F32, &device)?;

    // Should catch OOM and return error, NOT panic
    let result = tft.forward(&static_features, &historical_features, &future_features);

    if result.is_err() {
        // Verify error is Hardware category
        let err = result.unwrap_err();
        assert!(matches!(err, MLError::ModelError(_)));
    }

    drop(vram_consumer); // Clean up
    Ok(())
}

Test #13: Out-of-Distribution Inputs (Quantized)

File: ml/tests/tft_quantized_ood_test.rs (NEW)

#[test]
fn test_quantized_model_handles_out_of_distribution_inputs() -> Result<(), MLError> {
    let device = Device::Cpu;
    let config = TFTConfig::default();

    // Create quantized TFT
    let fp32_tft = TemporalFusionTransformer::new(config, device.clone())?;
    let quantized_tft = QuantizedTemporalFusionTransformer::from_fp32(&fp32_tft)?;

    // Create input with extreme value (calibration range typically -2.0 to 2.0)
    let mut static_data = vec![1.0f32; 5];
    static_data[2] = 50.0; // WAY outside calibration range
    let static_features = Tensor::from_vec(static_data, (1, 5), &device)?;

    let historical_features = Tensor::zeros((1, 60, 210), DType::F32, &device)?;
    let future_features = Tensor::zeros((1, 10, 10), DType::F32, &device)?;

    // Should NOT panic, should clamp to INT8 range
    let output = quantized_tft.forward(&static_features, &historical_features, &future_features)?;

    // Verify output is valid (no NaN)
    let output_vec = output.to_vec3::<f32>()?;
    let has_nan = output_vec.iter()
        .flat_map(|batch| batch.iter())
        .flat_map(|horizon| horizon.iter())
        .any(|&val| val.is_nan());
    assert!(!has_nan, "Quantized model should produce valid output for OOD inputs");
    Ok(())
}

Test #15: CUDA Initialization Failure Fallback

File: ml/tests/tft_cuda_fallback_test.rs (NEW)

#[test]
fn test_graceful_fallback_on_cuda_initialization_failure() -> Result<(), MLError> {
    // Simulate CUDA unavailable by using CPU device
    let device = Device::Cpu;
    let config = TFTConfig::default();

    // Should succeed on CPU
    let tft = TemporalFusionTransformer::new(config, device.clone())?;

    let static_features = Tensor::zeros((2, 5), DType::F32, &device)?;
    let historical_features = Tensor::zeros((2, 60, 210), DType::F32, &device)?;
    let future_features = Tensor::zeros((2, 10, 10), DType::F32, &device)?;

    let output = tft.forward(&static_features, &historical_features, &future_features)?;
    assert_eq!(output.dims(), &[2, 10, 9]);
    Ok(())
}

#[test]
#[ignore = "Requires CUDA_VISIBLE_DEVICES=-1 environment variable"]
fn test_cuda_fallback_with_env_variable() -> Result<(), MLError> {
    // Run with: CUDA_VISIBLE_DEVICES=-1 cargo test test_cuda_fallback_with_env_variable
    std::env::set_var("CUDA_VISIBLE_DEVICES", "-1");

    // Device::cuda_if_available should fall back to CPU
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    assert!(!device.is_cuda());

    let config = TFTConfig::default();
    let tft = TemporalFusionTransformer::new(config, device)?;

    // Should work fine
    assert!(true);
    Ok(())
}

Test #21: Feature Extraction NaN Handling

File: ml/tests/feature_extraction_nan_test.rs (NEW)

#[test]
fn test_feature_extraction_with_nan_ohlcv() -> Result<(), MLError> {
    use crate::features::extraction::extract_features;

    // Create OHLCV data with NaN
    let mut ohlcv = vec![
        (100.0, 105.0, 98.0, 103.0, 1000.0), // Valid
        (f32::NAN, 106.0, 99.0, 104.0, 1100.0), // NaN open
        (104.0, 107.0, 100.0, 105.0, 1200.0), // Valid
    ];

    // Should detect NaN and return error
    let result = extract_features(&ohlcv);
    assert!(result.is_err());

    match result.unwrap_err() {
        MLError::DataPreprocessing { stage, message } => {
            assert!(message.contains("NaN") || message.contains("invalid"));
        },
        _ => panic!("Expected DataPreprocessing error"),
    }
    Ok(())
}

---

Phase 2: High-Impact Edge Cases (8 hours)

Test #1: Zero Variance Normalization

File: ml/tests/cuda_compat_zero_variance_test.rs (NEW)

#[test]
fn test_layer_norm_with_zero_variance_input() -> Result<(), MLError> {
    let device = Device::Cpu;

    // Zero variance: all values identical
    let input = Tensor::new(&[[5.0f32, 5.0, 5.0], [2.0, 2.0, 2.0]], &device)?;
    let weight = Tensor::ones(3, DType::F32, &device)?;
    let bias = Tensor::zeros(3, DType::F32, &device)?;

    // Should NOT panic or produce NaN
    let output = cuda_layer_norm(&input, &[3], Some(&weight), Some(&bias), 1e-5)?;

    // Verify no NaN/Inf
    let output_vec = output.to_vec2::<f32>()?;
    for row in &output_vec {
        for &val in row {
            assert!(!val.is_nan(), "Output should not contain NaN");
            assert!(!val.is_infinite(), "Output should not contain Inf");
        }
    }

    // Output should be all zeros (x - mean = 0)
    for row in &output_vec {
        for &val in row {
            assert!(val.abs() < 1e-5, "Output should be near zero");
        }
    }
    Ok(())
}

Test #2: Device Mismatch Handling

File: Add to ml/tests/cuda_compat_device_test.rs (NEW)

#[test]
#[cfg(feature = "cuda")]
#[ignore = "Requires GPU"]
fn test_layer_norm_handles_device_mismatch() -> Result<(), MLError> {
    let cpu_device = Device::Cpu;
    let gpu_device = Device::cuda_if_available(0)?;

    if !gpu_device.is_cuda() {
        return Ok(()); // Skip if no GPU
    }

    // Input on GPU
    let input = Tensor::new(&[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]], &gpu_device)?;

    // Weight/bias on CPU (device mismatch)
    let weight = Tensor::ones(3, DType::F32, &cpu_device)?;
    let bias = Tensor::zeros(3, DType::F32, &cpu_device)?;

    // Should auto-move weight/bias to GPU, NOT panic
    let output = cuda_layer_norm(&input, &[3], Some(&weight), Some(&bias), 1e-5)?;

    assert_eq!(output.dims(), &[2, 3]);
    assert!(output.device().is_cuda());
    Ok(())
}

Test #7: LRU Cache Eviction

File: ml/tests/tft_lru_cache_test.rs (NEW)

#[test]
fn test_tft_state_lru_cache_eviction_under_load() -> Result<(), MLError> {
    use lru::LruCache;
    use std::num::NonZeroUsize;

    let capacity = NonZeroUsize::new(10).unwrap(); // Small cache for testing
    let mut cache: LruCache<String, Tensor> = LruCache::new(capacity);

    let device = Device::Cpu;

    // Insert 11 items (exceeds capacity of 10)
    for i in 0..11 {
        let key = format!("key_{}", i);
        let tensor = Tensor::zeros((2, 3), DType::F32, &device)?;
        cache.put(key, tensor);
    }

    // First item (key_0) should be evicted
    assert!(cache.get("key_0").is_none(), "LRU item should be evicted");

    // Last 10 items should remain
    for i in 1..11 {
        let key = format!("key_{}", i);
        assert!(cache.get(&key).is_some(), "Recent item should remain");
    }

    Ok(())
}

Test #18: AutoBatchSizer Zero GPU Memory

File: ml/tests/auto_batch_sizer_zero_gpu_test.rs (NEW)

#[test]
fn test_auto_batch_sizer_handles_zero_gpu_memory() -> Result<(), MLError> {
    use crate::memory_optimization::auto_batch_size::{AutoBatchSizer, BatchSizeConfig, OptimizerType};

    // Mock zero GPU memory scenario
    let config = BatchSizeConfig {
        model_memory_mb: 125.0,
        sequence_length: 60,
        feature_dim: 225,
        gradient_checkpointing: false,
        optimizer_type: OptimizerType::Adam,
        safety_margin: 0.20,
        available_memory_mb: 0.0, // ZERO available memory
    };

    // Should return minimum batch_size=1, NOT panic
    let batch_size = calculate_batch_size(&config)?;
    assert_eq!(batch_size, 1, "Should return minimum batch size");
    Ok(())
}

---

Phase 3: Robustness Improvements (10 hours)

Tests #3, #4, #10, #14, #16, #19, #20

(Templates similar to above, focusing on:)

• F64 CPU fallback
• Sigmoid extreme values
• Shared VarMap guard
• Zero variance quantization
• Thread-safe metrics
• Batch > data size
• Single sample DBN

---

Phase 4: Validation (2 hours)

Tests #12, #17

(Confirm existing guardrails work)

---

Test File Organization

ml/tests/
├── tft_zero_batch_test.rs              # Test #5
├── tft_batch_mismatch_test.rs          # Test #6
├── tft_nan_handling_test.rs            # Test #8
├── tft_corrupt_checkpoint_test.rs      # Test #9
├── tft_gpu_oom_test.rs                 # Test #11
├── tft_quantized_ood_test.rs           # Test #13
├── tft_cuda_fallback_test.rs           # Test #15
├── feature_extraction_nan_test.rs      # Test #21
├── cuda_compat_zero_variance_test.rs   # Test #1
├── cuda_compat_device_test.rs          # Test #2
├── tft_lru_cache_test.rs               # Test #7
├── auto_batch_sizer_zero_gpu_test.rs   # Test #18
└── ...

---

Running Tests

# Run all new tests
cargo test -p ml --test tft_zero_batch_test
cargo test -p ml --test tft_batch_mismatch_test
cargo test -p ml --test tft_nan_handling_test

# Run GPU-specific tests (requires GPU)
cargo test -p ml --test tft_gpu_oom_test --features cuda -- --ignored

# Run all ml tests
cargo test -p ml

# Run with verbose output
cargo test -p ml -- --nocapture

---

Common Test Utilities

Add to ml/tests/common/mod.rs:

pub fn create_test_tft(device: Device) -> Result<TemporalFusionTransformer, MLError> {
    let config = TFTConfig::default();
    TemporalFusionTransformer::new(config, device)
}

pub fn create_zero_batch_inputs(device: &Device) -> Result<(Tensor, Tensor, Tensor), MLError> {
    let static_features = Tensor::zeros((0, 5), DType::F32, device)?;
    let historical_features = Tensor::zeros((0, 60, 210), DType::F32, device)?;
    let future_features = Tensor::zeros((0, 10, 10), DType::F32, device)?;
    Ok((static_features, historical_features, future_features))
}

pub fn assert_no_nan_inf(tensor: &Tensor) -> Result<(), MLError> {
    let vec = tensor.flatten_all()?.to_vec1::<f32>()?;
    for &val in &vec {
        assert!(!val.is_nan(), "Tensor contains NaN");
        assert!(!val.is_infinite(), "Tensor contains Inf");
    }
    Ok(())
}

---

Time Estimates (Conservative)

Phase	Tests	Hours/Test	Total
Phase 1	8	2.0	16h
Phase 2	4	2.0	8h
Phase 3	7	1.5	10.5h
Phase 4	2	1.0	2h
TOTAL	21	1.75 avg	36.5h

Realistic Timeline (1 developer, part-time):

• Week 1: Phase 1 (8 tests, 16h)
• Week 2: Phase 2 (4 tests, 8h)
• Week 3: Phase 3 (7 tests, 10.5h)
• Week 4: Phase 4 (2 tests, 2h)

---

Success Criteria

• ✅ All 21 tests pass on first run
• ✅ No test introduces new unwrap/expect calls
• ✅ GPU tests pass on CUDA-enabled systems
• ✅ CPU fallback tests pass on non-GPU systems
• ✅ Test execution time < 5 minutes total (excluding GPU OOM test)
• ✅ Code coverage increases by ≥5%

---

Generated by: Agent 23 Date: 2025-10-25 Next Step: Implement Phase 1 tests before Runpod deployment