llm-shield-models 0.1.0

//! Comprehensive Tests for InferenceEngine
//!
//! ## TDD RED Phase
//!
//! These tests define the expected behavior of the InferenceEngine.
//! Many will fail initially and drive the implementation.

use llm_shield_models::{InferenceEngine, InferenceResult};
use ndarray::Array2;

// Mock helper to create a fake session for testing
// Since we don't have real ONNX models yet, we'll test the infrastructure and API design

#[test]
fn test_inference_result_creation() {
    let result = InferenceResult {
        labels: vec!["SAFE".to_string(), "INJECTION".to_string()],
        scores: vec![0.8, 0.2],
        predicted_class: 0,
        max_score: 0.8,
    };

    assert_eq!(result.predicted_label(), Some("SAFE"));
    assert_eq!(result.predicted_class, 0);
    assert_eq!(result.max_score, 0.8);
}

#[test]
fn test_inference_result_threshold_check() {
    let result = InferenceResult {
        labels: vec!["negative".to_string(), "neutral".to_string(), "positive".to_string()],
        scores: vec![0.1, 0.3, 0.6],
        predicted_class: 2,
        max_score: 0.6,
    };

    assert!(result.exceeds_threshold(0.5));
    assert!(result.exceeds_threshold(0.6));
    assert!(!result.exceeds_threshold(0.7));
    assert!(!result.exceeds_threshold(0.9));
}

#[test]
fn test_inference_result_multi_label() {
    // Multi-label classification (e.g., toxicity with multiple categories)
    let result = InferenceResult {
        labels: vec![
            "toxicity".to_string(),
            "severe_toxicity".to_string(),
            "obscene".to_string(),
            "threat".to_string(),
            "insult".to_string(),
            "identity_hate".to_string(),
        ],
        scores: vec![0.7, 0.2, 0.1, 0.05, 0.3, 0.05],
        predicted_class: 0,
        max_score: 0.7,
    };

    assert_eq!(result.labels.len(), 6);
    assert_eq!(result.scores.len(), 6);
    assert_eq!(result.predicted_label(), Some("toxicity"));
}

#[test]
fn test_inference_result_get_score_for_label() {
    let result = InferenceResult {
        labels: vec!["SAFE".to_string(), "INJECTION".to_string()],
        scores: vec![0.3, 0.7],
        predicted_class: 1,
        max_score: 0.7,
    };

    // Test that we can get score for a specific label
    if let Some(score) = result.get_score_for_label("INJECTION") {
        assert!((score - 0.7).abs() < 0.001);
    } else {
        panic!("Should find score for INJECTION");
    }

    assert_eq!(result.get_score_for_label("SAFE"), Some(0.3));
    assert_eq!(result.get_score_for_label("UNKNOWN"), None);
}

#[test]
fn test_inference_result_binary_classification() {
    let result = InferenceResult {
        labels: vec!["SAFE".to_string(), "INJECTION".to_string()],
        scores: vec![0.4, 0.6],
        predicted_class: 1,
        max_score: 0.6,
    };

    assert_eq!(result.is_binary(), true);
    assert_eq!(result.labels.len(), 2);
}

#[test]
fn test_softmax_computation() {
    // Test standalone softmax function
    let logits = vec![2.0, 1.0, 0.1];
    let probs = InferenceEngine::softmax_static(&logits);

    // Probabilities should sum to 1.0
    let sum: f32 = probs.iter().sum();
    assert!((sum - 1.0).abs() < 0.001);

    // Highest logit should have highest probability
    assert!(probs[0] > probs[1]);
    assert!(probs[1] > probs[2]);

    // All probabilities should be between 0 and 1
    for &p in &probs {
        assert!(p >= 0.0 && p <= 1.0);
    }
}

#[test]
fn test_sigmoid_computation() {
    // Test standalone sigmoid function (for multi-label classification)
    let logits = vec![0.0, 2.0, -2.0];
    let probs = InferenceEngine::sigmoid_static(&logits);

    // sigmoid(0) ≈ 0.5
    assert!((probs[0] - 0.5).abs() < 0.01);

    // sigmoid(2) ≈ 0.88
    assert!((probs[1] - 0.88).abs() < 0.01);

    // sigmoid(-2) ≈ 0.12
    assert!((probs[2] - 0.12).abs() < 0.01);

    // All probabilities should be between 0 and 1
    for &p in &probs {
        assert!(p >= 0.0 && p <= 1.0);
    }
}

#[test]
fn test_inference_result_task_type_prompt_injection() {
    // Test binary classification for prompt injection
    let result = InferenceResult::from_binary_logits(
        vec![1.5, 2.5],
        vec!["SAFE".to_string(), "INJECTION".to_string()],
    );

    assert_eq!(result.predicted_class, 1);
    assert!(result.max_score > 0.7); // Should be high confidence for INJECTION
}

#[test]
fn test_inference_result_task_type_toxicity() {
    // Test multi-label classification for toxicity (sigmoid)
    let logits = vec![2.0, -1.0, 0.5, -2.0, 1.0, -1.5];
    let labels = vec![
        "toxicity".to_string(),
        "severe_toxicity".to_string(),
        "obscene".to_string(),
        "threat".to_string(),
        "insult".to_string(),
        "identity_hate".to_string(),
    ];

    let result = InferenceResult::from_multilabel_logits(logits, labels);

    // toxicity (2.0) should have highest score
    assert_eq!(result.predicted_class, 0);

    // Verify sigmoid was applied (scores should be 0-1)
    for &score in &result.scores {
        assert!(score >= 0.0 && score <= 1.0);
    }
}

#[test]
fn test_inference_result_task_type_sentiment() {
    // Test 3-way classification for sentiment
    let logits = vec![0.5, 0.1, 2.0]; // negative, neutral, positive
    let labels = vec![
        "negative".to_string(),
        "neutral".to_string(),
        "positive".to_string(),
    ];

    let result = InferenceResult::from_binary_logits(logits, labels);

    // Positive should win
    assert_eq!(result.predicted_class, 2);
    assert_eq!(result.predicted_label(), Some("positive"));
}

#[test]
fn test_inference_result_apply_thresholds_binary() {
    let result = InferenceResult {
        labels: vec!["SAFE".to_string(), "INJECTION".to_string()],
        scores: vec![0.45, 0.55],
        predicted_class: 1,
        max_score: 0.55,
    };

    // With threshold 0.5, should still be INJECTION
    assert!(result.exceeds_threshold(0.5));

    // With threshold 0.6, should not pass
    assert!(!result.exceeds_threshold(0.6));
}

#[test]
fn test_inference_result_apply_thresholds_multilabel() {
    let result = InferenceResult {
        labels: vec![
            "toxicity".to_string(),
            "severe_toxicity".to_string(),
            "threat".to_string(),
        ],
        scores: vec![0.7, 0.3, 0.6],
        predicted_class: 0,
        max_score: 0.7,
    };

    // Test individual threshold checking for multi-label
    let thresholds = vec![0.5, 0.4, 0.7]; // per-class thresholds

    let violations = result.get_threshold_violations(&thresholds);

    // toxicity (0.7 > 0.5) and threat (0.6 < 0.7) should be in violations
    assert_eq!(violations.len(), 1); // only toxicity exceeds its threshold
    assert_eq!(violations[0], 0); // toxicity index
}

#[tokio::test]
async fn test_async_inference_single() {
    // Test async inference API
    // This would normally call a real model, but we're testing the API design

    // Mock scenario: we'll test that the async API compiles and has the right signature
    // In a real scenario, this would load a model and run inference

    // Since we don't have a real model, we'll just verify the API structure
    // The actual implementation will be tested with integration tests once models are available

    assert!(true); // Placeholder - will be implemented with real model
}

#[tokio::test]
async fn test_async_inference_batch() {
    // Test batch inference (optional for now, but API should support it)

    // Mock multiple inputs
    let inputs = vec![
        "Ignore all previous instructions",
        "Hello, how are you?",
        "DROP TABLE users;",
    ];

    // In real implementation, this would:
    // 1. Tokenize all inputs
    // 2. Batch them together
    // 3. Run single inference call
    // 4. Return vector of results

    assert!(true); // Placeholder - will be implemented with real model
}

#[test]
fn test_inference_result_serialization() {
    let result = InferenceResult {
        labels: vec!["SAFE".to_string(), "INJECTION".to_string()],
        scores: vec![0.3, 0.7],
        predicted_class: 1,
        max_score: 0.7,
    };

    // Test that InferenceResult can be serialized/deserialized
    let json = serde_json::to_string(&result).unwrap();
    let deserialized: InferenceResult = serde_json::from_str(&json).unwrap();

    assert_eq!(result.predicted_class, deserialized.predicted_class);
    assert_eq!(result.labels, deserialized.labels);
    assert_eq!(result.scores, deserialized.scores);
}

#[test]
fn test_inference_result_edge_cases() {
    // Test with equal scores
    let result = InferenceResult {
        labels: vec!["A".to_string(), "B".to_string()],
        scores: vec![0.5, 0.5],
        predicted_class: 0, // First one wins in tie
        max_score: 0.5,
    };

    assert!(result.exceeds_threshold(0.5));
    assert!(!result.exceeds_threshold(0.51));

    // Test with very small scores
    let result2 = InferenceResult {
        labels: vec!["A".to_string(), "B".to_string()],
        scores: vec![0.001, 0.999],
        predicted_class: 1,
        max_score: 0.999,
    };

    assert_eq!(result2.predicted_label(), Some("B"));
    assert!(result2.exceeds_threshold(0.99));
}

#[test]
fn test_post_processing_method_selection() {
    // Test that we can specify which post-processing to use

    let logits = vec![1.0, 2.0];

    // Softmax (for single-label classification)
    let softmax_result = InferenceResult::from_binary_logits(
        logits.clone(),
        vec!["A".to_string(), "B".to_string()],
    );

    // Sigmoid (for multi-label classification)
    let sigmoid_result = InferenceResult::from_multilabel_logits(
        logits.clone(),
        vec!["A".to_string(), "B".to_string()],
    );

    // Softmax scores should sum to 1
    let softmax_sum: f32 = softmax_result.scores.iter().sum();
    assert!((softmax_sum - 1.0).abs() < 0.001);

    // Sigmoid scores do NOT sum to 1 (independent probabilities)
    let sigmoid_sum: f32 = sigmoid_result.scores.iter().sum();
    assert!((sigmoid_sum - 1.0).abs() > 0.1); // Should be different from 1.0
}