foxhunt/storage/tests/model_helpers_tests.rs

//! Comprehensive tests for model_helpers module
//!
//! Tests connection pooling, retry logic, model info/version structures,
//! and helper functions for model storage operations.

use std::collections::HashMap;
use std::sync::Arc;

use chrono::Utc;
use object_store::memory::InMemory;
use storage::model_helpers::{
    ConnectionPool, ModelInfo, ModelVersion, ProgressCallback, RetryConfig, TrainingInfo,
};

#[test]
fn test_model_version_creation() {
    let version = ModelVersion {
        name: "mamba".to_string(),
        version: "v1.0".to_string(),
        path: "models/mamba/v1.0/weights.safetensors".to_string(),
        size: 1024 * 1024,
        created_at: Utc::now(),
        metrics: HashMap::new(),
        training_info: None,
        checksum: Some("abc123".to_string()),
    };

    assert_eq!(version.name, "mamba");
    assert_eq!(version.version, "v1.0");
    assert_eq!(version.size, 1024 * 1024);
    assert!(version.checksum.is_some());
}

#[test]
fn test_model_version_with_metrics() {
    let mut metrics = HashMap::new();
    metrics.insert("accuracy".to_string(), 0.95);
    metrics.insert("loss".to_string(), 0.05);

    let version = ModelVersion {
        name: "dqn".to_string(),
        version: "v2.1".to_string(),
        path: "models/dqn/v2.1/checkpoint.pt".to_string(),
        size: 2048 * 1024,
        created_at: Utc::now(),
        metrics: metrics.clone(),
        training_info: None,
        checksum: None,
    };

    assert_eq!(version.metrics.len(), 2);
    assert_eq!(version.metrics.get("accuracy"), Some(&0.95));
    assert_eq!(version.metrics.get("loss"), Some(&0.05));
}

#[test]
fn test_training_info_creation() {
    let training_info = TrainingInfo {
        epoch: 100,
        step: 50000,
        validation_loss: Some(0.03),
        training_loss: Some(0.02),
        duration_seconds: 3600,
        git_commit: Some("abc123def456".to_string()),
    };

    assert_eq!(training_info.epoch, 100);
    assert_eq!(training_info.step, 50000);
    assert_eq!(training_info.validation_loss, Some(0.03));
    assert_eq!(training_info.training_loss, Some(0.02));
    assert_eq!(training_info.duration_seconds, 3600);
    assert!(training_info.git_commit.is_some());
}

#[test]
fn test_model_version_with_training_info() {
    let training_info = TrainingInfo {
        epoch: 50,
        step: 25000,
        validation_loss: Some(0.04),
        training_loss: Some(0.03),
        duration_seconds: 1800,
        git_commit: None,
    };

    let version = ModelVersion {
        name: "ppo".to_string(),
        version: "v1.5".to_string(),
        path: "models/ppo/v1.5/model.safetensors".to_string(),
        size: 512 * 1024,
        created_at: Utc::now(),
        metrics: HashMap::new(),
        training_info: Some(training_info),
        checksum: None,
    };

    assert!(version.training_info.is_some());
    let info = version.training_info.unwrap();
    assert_eq!(info.epoch, 50);
    assert_eq!(info.step, 25000);
}

#[test]
fn test_model_info_creation() {
    let version1 = ModelVersion {
        name: "tft".to_string(),
        version: "v1.0".to_string(),
        path: "models/tft/v1.0/weights.bin".to_string(),
        size: 1024,
        created_at: Utc::now(),
        metrics: HashMap::new(),
        training_info: None,
        checksum: None,
    };

    let version2 = ModelVersion {
        name: "tft".to_string(),
        version: "v2.0".to_string(),
        path: "models/tft/v2.0/weights.bin".to_string(),
        size: 2048,
        created_at: Utc::now(),
        metrics: HashMap::new(),
        training_info: None,
        checksum: None,
    };

    let mut tags = HashMap::new();
    tags.insert("framework".to_string(), "candle".to_string());

    let model_info = ModelInfo {
        name: "tft".to_string(),
        versions: vec![version1.clone(), version2.clone()],
        total_size: 3072,
        latest_version: Some(version2),
        architecture: Some("transformer".to_string()),
        last_updated: Utc::now(),
        tags,
    };

    assert_eq!(model_info.name, "tft");
    assert_eq!(model_info.versions.len(), 2);
    assert_eq!(model_info.total_size, 3072);
    assert!(model_info.latest_version.is_some());
    assert_eq!(model_info.architecture, Some("transformer".to_string()));
    assert_eq!(model_info.tags.len(), 1);
}

#[test]
fn test_model_info_no_versions() {
    let model_info = ModelInfo {
        name: "empty_model".to_string(),
        versions: vec![],
        total_size: 0,
        latest_version: None,
        architecture: None,
        last_updated: Utc::now(),
        tags: HashMap::new(),
    };

    assert_eq!(model_info.versions.len(), 0);
    assert_eq!(model_info.total_size, 0);
    assert!(model_info.latest_version.is_none());
}

#[test]
fn test_retry_config_default() {
    let config = RetryConfig::default();

    assert_eq!(config.max_attempts, 3);
    assert_eq!(
        config.initial_delay,
        std::time::Duration::from_millis(100)
    );
    assert_eq!(config.max_delay, std::time::Duration::from_secs(10));
    assert_eq!(config.backoff_multiplier, 2.0);
}

#[test]
fn test_retry_config_custom() {
    let config = RetryConfig {
        max_attempts: 5,
        initial_delay: std::time::Duration::from_millis(50),
        max_delay: std::time::Duration::from_secs(5),
        backoff_multiplier: 1.5,
    };

    assert_eq!(config.max_attempts, 5);
    assert_eq!(config.initial_delay, std::time::Duration::from_millis(50));
    assert_eq!(config.max_delay, std::time::Duration::from_secs(5));
    assert_eq!(config.backoff_multiplier, 1.5);
}

#[tokio::test]
async fn test_connection_pool_creation() {
    let store1 = Arc::new(InMemory::new());
    let store2 = Arc::new(InMemory::new());

    let pool = ConnectionPool::new(vec![store1, store2]);

    // Get store should succeed
    let store = pool.get_store().await;
    assert!(store.is_ok());
}

#[tokio::test]
async fn test_connection_pool_empty() {
    let pool = ConnectionPool::new(vec![]);

    // Empty pool should fail
    let result = pool.get_store().await;
    assert!(result.is_err());
}

#[tokio::test]
async fn test_connection_pool_round_robin() {
    let store1 = Arc::new(InMemory::new());
    let store2 = Arc::new(InMemory::new());
    let store3 = Arc::new(InMemory::new());

    let pool = ConnectionPool::new(vec![store1, store2, store3]);

    // Get multiple stores to verify round-robin
    let _s1 = pool.get_store().await.unwrap();
    let _s2 = pool.get_store().await.unwrap();
    let _s3 = pool.get_store().await.unwrap();
    // Should wrap around
    let _s4 = pool.get_store().await.unwrap();
}

#[tokio::test]
async fn test_connection_pool_concurrent_access() {
    let store1 = Arc::new(InMemory::new());
    let store2 = Arc::new(InMemory::new());

    let pool = Arc::new(ConnectionPool::new(vec![store1, store2]));

    let mut handles = vec![];

    for _ in 0..10 {
        let pool = pool.clone();
        let handle = tokio::spawn(async move {
            let store = pool.get_store().await.unwrap();
            // Use the store (just verify we got one)
            drop(store);
        });
        handles.push(handle);
    }

    for handle in handles {
        handle.await.unwrap();
    }
}

#[test]
fn test_progress_callback_creation() {
    let callback: ProgressCallback = Arc::new(|downloaded, total| {
        assert!(downloaded <= total);
    });

    // Call the callback
    callback(50, 100);
    callback(100, 100);
}

#[test]
fn test_progress_callback_with_state() {
    let progress_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
    let progress_count_clone = progress_count.clone();

    let callback: ProgressCallback = Arc::new(move |downloaded, total| {
        progress_count_clone.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        assert!(downloaded <= total);
    });

    // Call multiple times
    callback(25, 100);
    callback(50, 100);
    callback(100, 100);

    assert_eq!(
        progress_count.load(std::sync::atomic::Ordering::SeqCst),
        3
    );
}

#[test]
fn test_model_version_serialization() {
    let version = ModelVersion {
        name: "test_model".to_string(),
        version: "v1.0".to_string(),
        path: "models/test/v1.0/weights.bin".to_string(),
        size: 1024,
        created_at: Utc::now(),
        metrics: HashMap::new(),
        training_info: None,
        checksum: Some("hash123".to_string()),
    };

    // Serialize to JSON
    let json = serde_json::to_string(&version).expect("serialization failed");
    assert!(json.contains("test_model"));
    assert!(json.contains("v1.0"));

    // Deserialize back
    let deserialized: ModelVersion = serde_json::from_str(&json).expect("deserialization failed");
    assert_eq!(deserialized.name, version.name);
    assert_eq!(deserialized.version, version.version);
}

#[test]
fn test_model_info_serialization() {
    let mut tags = HashMap::new();
    tags.insert("env".to_string(), "production".to_string());

    let model_info = ModelInfo {
        name: "test_model".to_string(),
        versions: vec![],
        total_size: 1024,
        latest_version: None,
        architecture: Some("transformer".to_string()),
        last_updated: Utc::now(),
        tags,
    };

    // Serialize
    let json = serde_json::to_string(&model_info).expect("serialization failed");
    assert!(json.contains("test_model"));
    assert!(json.contains("transformer"));
    assert!(json.contains("production"));

    // Deserialize
    let deserialized: ModelInfo = serde_json::from_str(&json).expect("deserialization failed");
    assert_eq!(deserialized.name, model_info.name);
    assert_eq!(deserialized.architecture, model_info.architecture);
}

#[test]
fn test_training_info_serialization() {
    let training_info = TrainingInfo {
        epoch: 100,
        step: 50000,
        validation_loss: Some(0.03),
        training_loss: Some(0.02),
        duration_seconds: 3600,
        git_commit: Some("abc123".to_string()),
    };

    let json = serde_json::to_string(&training_info).expect("serialization failed");
    assert!(json.contains("100")); // epoch
    assert!(json.contains("50000")); // step

    let deserialized: TrainingInfo = serde_json::from_str(&json).expect("deserialization failed");
    assert_eq!(deserialized.epoch, training_info.epoch);
    assert_eq!(deserialized.step, training_info.step);
}

#[test]
fn test_model_version_with_large_metrics() {
    let mut metrics = HashMap::new();
    for i in 0..100 {
        metrics.insert(format!("metric_{}", i), i as f64);
    }

    let version = ModelVersion {
        name: "large_model".to_string(),
        version: "v1.0".to_string(),
        path: "models/large/v1.0/weights.bin".to_string(),
        size: 1024 * 1024,
        created_at: Utc::now(),
        metrics,
        training_info: None,
        checksum: None,
    };

    assert_eq!(version.metrics.len(), 100);
    assert_eq!(version.metrics.get("metric_0"), Some(&0.0));
    assert_eq!(version.metrics.get("metric_99"), Some(&99.0));
}

#[test]
fn test_model_info_with_many_versions() {
    let mut versions = vec![];
    for i in 0..50 {
        versions.push(ModelVersion {
            name: "multi_version_model".to_string(),
            version: format!("v{}.0", i),
            path: format!("models/multi/v{}.0/weights.bin", i),
            size: 1024 * (i + 1),
            created_at: Utc::now(),
            metrics: HashMap::new(),
            training_info: None,
            checksum: None,
        });
    }

    let total_size: u64 = versions.iter().map(|v| v.size).sum();

    let model_info = ModelInfo {
        name: "multi_version_model".to_string(),
        versions: versions.clone(),
        total_size,
        latest_version: versions.last().cloned(),
        architecture: None,
        last_updated: Utc::now(),
        tags: HashMap::new(),
    };

    assert_eq!(model_info.versions.len(), 50);
    assert!(model_info.latest_version.is_some());
    assert_eq!(model_info.latest_version.unwrap().version, "v49.0");
}

#[test]
fn test_training_info_optional_fields() {
    // All optional fields None
    let info1 = TrainingInfo {
        epoch: 10,
        step: 1000,
        validation_loss: None,
        training_loss: None,
        duration_seconds: 600,
        git_commit: None,
    };

    assert!(info1.validation_loss.is_none());
    assert!(info1.training_loss.is_none());
    assert!(info1.git_commit.is_none());

    // All optional fields Some
    let info2 = TrainingInfo {
        epoch: 20,
        step: 2000,
        validation_loss: Some(0.01),
        training_loss: Some(0.005),
        duration_seconds: 1200,
        git_commit: Some("def456".to_string()),
    };

    assert_eq!(info2.validation_loss, Some(0.01));
    assert_eq!(info2.training_loss, Some(0.005));
    assert_eq!(info2.git_commit, Some("def456".to_string()));
}

#[test]
fn test_model_info_tags() {
    let mut tags = HashMap::new();
    tags.insert("env".to_string(), "staging".to_string());
    tags.insert("version".to_string(), "beta".to_string());
    tags.insert("team".to_string(), "ml_team".to_string());

    let model_info = ModelInfo {
        name: "tagged_model".to_string(),
        versions: vec![],
        total_size: 0,
        latest_version: None,
        architecture: None,
        last_updated: Utc::now(),
        tags: tags.clone(),
    };

    assert_eq!(model_info.tags.len(), 3);
    assert_eq!(model_info.tags.get("env"), Some(&"staging".to_string()));
    assert_eq!(model_info.tags.get("version"), Some(&"beta".to_string()));
    assert_eq!(model_info.tags.get("team"), Some(&"ml_team".to_string()));
}