foxhunt/crates/ml/tests/tft_lru_cache_test.rs

/// Test to verify TFT attention cache LRU behavior
///
/// This test ensures the unbounded HashMap memory leak fix (2025-10-25)
/// works correctly by validating:
/// - Cache size never exceeds MAX_CACHE_ENTRIES (2000)
/// - LRU eviction happens automatically
/// - Memory is bounded even with many insertions
use anyhow::Result;
use ml::tft::{TFTConfig, TFTState};

#[test]
fn test_tft_state_lru_cache_bounds() -> Result<()> {
    let config = TFTConfig::default();
    let mut state = TFTState::zeros(&config)?;

    // Verify initial state
    assert_eq!(state.attention_cache.len(), 0, "Cache should start empty");

    // Insert 3000 entries (exceeds MAX_CACHE_ENTRIES of 2000)
    for i in 0..3000 {
        let key = format!("cache_key_{}", i);
        let value = candle_core::Tensor::zeros(
            &[8, 64],
            candle_core::DType::F32,
            &candle_core::Device::Cpu,
        )?;
        state.attention_cache.put(key, value);
    }

    // Cache should be capped at MAX_CACHE_ENTRIES (2000)
    assert_eq!(
        state.attention_cache.len(),
        TFTState::MAX_CACHE_ENTRIES,
        "Cache should never exceed MAX_CACHE_ENTRIES (2000)"
    );

    // Oldest entries (0-999) should be evicted by LRU policy
    assert!(
        state.attention_cache.get("cache_key_0").is_none(),
        "Oldest entry should be evicted"
    );
    assert!(
        state.attention_cache.get("cache_key_999").is_none(),
        "Old entries should be evicted"
    );

    // Newest entries (1000-2999) should still be present
    assert!(
        state.attention_cache.get("cache_key_1000").is_some(),
        "Recent entry should be retained"
    );
    assert!(
        state.attention_cache.get("cache_key_2999").is_some(),
        "Newest entry should be retained"
    );

    Ok(())
}

#[test]
fn test_tft_state_cache_max_entries_constant() {
    // Verify MAX_CACHE_ENTRIES is sensible
    assert_eq!(
        TFTState::MAX_CACHE_ENTRIES,
        2000,
        "MAX_CACHE_ENTRIES should be 2000 (chosen for ~48MB overhead, 60% speedup)"
    );
}

#[test]
fn test_tft_state_creation_with_lru() -> Result<()> {
    let config = TFTConfig::default();
    let state = TFTState::zeros(&config)?;

    // Verify state is created with correct initial values
    assert!(state.hidden_state.is_none(), "Hidden state should be None");
    assert_eq!(state.last_update, 0, "Last update should be 0");
    assert_eq!(state.attention_cache.len(), 0, "Cache should be empty");

    Ok(())
}

#[test]
fn test_lru_eviction_order() -> Result<()> {
    let config = TFTConfig::default();
    let mut state = TFTState::zeros(&config)?;

    // Insert exactly MAX_CACHE_ENTRIES (1000) items
    for i in 0..TFTState::MAX_CACHE_ENTRIES {
        let key = format!("key_{}", i);
        let value = candle_core::Tensor::zeros(
            &[4, 32],
            candle_core::DType::F32,
            &candle_core::Device::Cpu,
        )?;
        state.attention_cache.put(key, value);
    }

    assert_eq!(state.attention_cache.len(), TFTState::MAX_CACHE_ENTRIES);

    // Access key_500 to make it recently used
    let _ = state.attention_cache.get("key_500");

    // Insert one more item (should evict key_0, the least recently used)
    state.attention_cache.put(
        "new_key".to_string(),
        candle_core::Tensor::zeros(&[4, 32], candle_core::DType::F32, &candle_core::Device::Cpu)?,
    );

    // key_0 should be evicted (oldest)
    assert!(
        state.attention_cache.get("key_0").is_none(),
        "key_0 should be evicted as LRU"
    );

    // key_500 should still be present (was accessed recently)
    assert!(
        state.attention_cache.get("key_500").is_some(),
        "key_500 should remain (recently accessed)"
    );

    // new_key should be present
    assert!(
        state.attention_cache.get("new_key").is_some(),
        "new_key should be present"
    );

    Ok(())
}