foxhunt/testing/integration/integration/mamba2_integration.rs

//! Mamba2 SSM model integration test
//!
//! Verifies: sequence buffering -> SSM forward pass -> sigmoid direction/confidence
//! Mamba2 requires sequence_length observations before producing real predictions.
//! Uses lightweight config for fast execution (<10s)

use ml::ensemble::adapters::Mamba2InferenceAdapter;
use ml::ensemble::inference_adapter::{FeatureVector, ModelInferenceAdapter};
use ml::mamba::Mamba2Config;

fn small_mamba2_config() -> Mamba2Config {
    Mamba2Config {
        d_model: 32,
        d_state: 8,
        d_head: 8,
        num_heads: 2,
        expand: 2,
        num_layers: 1,
        max_seq_len: 8,
        dropout: 0.0,
        ..Default::default()
    }
}

const SEQ_LEN: usize = 4;

#[test]
fn test_mamba2_adapter_buffers_before_predicting() {
    let adapter = Mamba2InferenceAdapter::new(small_mamba2_config(), SEQ_LEN)
        .expect("Mamba2InferenceAdapter::new should succeed");

    assert_eq!(adapter.model_name(), "MAMBA-2");
    assert!(
        !adapter.is_ready(),
        "Mamba2 should not be ready with empty buffer"
    );

    // Feed SEQ_LEN - 1 feature vectors: should return neutral
    for i in 0..(SEQ_LEN - 1) {
        let fv = FeatureVector {
            values: vec![0.1 * (i as f64 + 1.0); 51],
            timestamp: 1_700_000_000 + i as i64,
        };
        let pred = adapter.predict(&fv).expect("predict should not error");
        assert_eq!(pred.direction, 0.0, "Neutral direction while buffering");
        assert_eq!(pred.confidence, 0.0, "Zero confidence while buffering");
    }
}

#[test]
fn test_mamba2_produces_valid_prediction_after_warmup() {
    let adapter = Mamba2InferenceAdapter::new(small_mamba2_config(), SEQ_LEN)
        .expect("Mamba2InferenceAdapter::new should succeed");

    // Fill buffer
    for i in 0..SEQ_LEN {
        let fv = FeatureVector {
            values: vec![0.1; 51],
            timestamp: 1_700_000_000 + i as i64,
        };
        let _ = adapter.predict(&fv);
    }

    assert!(
        adapter.is_ready(),
        "Mamba2 should be ready after filling buffer"
    );

    // Real prediction
    let fv = FeatureVector {
        values: vec![0.2; 51],
        timestamp: 1_700_000_000 + SEQ_LEN as i64,
    };
    let pred = adapter.predict(&fv).expect("Mamba2 predict should succeed");

    assert!(
        pred.direction >= -1.0 && pred.direction <= 1.0,
        "direction {} out of [-1,1]",
        pred.direction
    );
    assert!(
        pred.confidence >= 0.0 && pred.confidence <= 1.0,
        "confidence {} out of [0,1]",
        pred.confidence
    );
    assert!(pred.direction.is_finite(), "direction must not be NaN/Inf");
    assert!(
        pred.confidence.is_finite(),
        "confidence must not be NaN/Inf"
    );
}

#[test]
fn test_mamba2_deterministic_with_same_sequence() {
    let adapter = Mamba2InferenceAdapter::new(small_mamba2_config(), SEQ_LEN)
        .expect("Mamba2InferenceAdapter::new should succeed");

    // Fill buffer with identical values
    for _ in 0..SEQ_LEN {
        let fv = FeatureVector {
            values: vec![0.1; 51],
            timestamp: 1_700_000_000,
        };
        let _ = adapter.predict(&fv);
    }

    // Two predictions with same input should be deterministic
    let fv = FeatureVector {
        values: vec![0.1; 51],
        timestamp: 1_700_000_000,
    };
    let pred1 = adapter.predict(&fv).expect("predict1");
    let pred2 = adapter.predict(&fv).expect("predict2");

    assert!(
        (pred1.direction - pred2.direction).abs() < 1e-6,
        "Mamba2 should be deterministic: {} vs {}",
        pred1.direction,
        pred2.direction
    );
}