foxhunt/ml/tests/checkpoint_test.rs

//! Checkpoint and Model Persistence Tests
//!
//! Comprehensive testing for model checkpointing covering:
//! - Checkpoint format validation
//! - Compression type selection
//! - Metadata creation and validation
//! - Storage backend operations
//! - Versioning and compatibility

#![allow(unused_crate_dependencies)]

use ml::checkpoint::{CheckpointFormat, CheckpointMetadata, CompressionType, ModelType};

/// Test: Checkpoint format variants
#[test]
fn test_checkpoint_format_variants() {
    // Test all format types exist
    let binary = CheckpointFormat::Binary;
    let json = CheckpointFormat::JSON;
    let msgpack = CheckpointFormat::MessagePack;
    let custom = CheckpointFormat::Custom;

    // Verify equality
    assert_eq!(binary, CheckpointFormat::Binary);
    assert_eq!(json, CheckpointFormat::JSON);
    assert_eq!(msgpack, CheckpointFormat::MessagePack);
    assert_eq!(custom, CheckpointFormat::Custom);
}

/// Test: Checkpoint format - Binary is fastest
#[test]
fn test_checkpoint_format_binary_performance() {
    let format = CheckpointFormat::Binary;

    // Binary format should be the default for performance
    assert_eq!(format, CheckpointFormat::Binary);
}

/// Test: Checkpoint format - JSON is human-readable
#[test]
fn test_checkpoint_format_json_readable() {
    let format = CheckpointFormat::JSON;

    // JSON is for debugging and inspection
    assert_eq!(format, CheckpointFormat::JSON);
}

/// Test: Checkpoint format - serialization
#[test]
fn test_checkpoint_format_serialization() {
    let format = CheckpointFormat::Binary;

    // Serialize to JSON
    let json = serde_json::to_string(&format).expect("Should serialize");

    // Deserialize back
    let deserialized: CheckpointFormat =
        serde_json::from_str(&json).expect("Should deserialize");

    assert_eq!(format, deserialized);
}

/// Test: Compression type variants
#[test]
fn test_compression_type_variants() {
    // Test all compression types
    let none = CompressionType::None;
    let lz4 = CompressionType::LZ4;
    let zstd = CompressionType::Zstd;
    let gzip = CompressionType::Gzip;

    // Verify equality
    assert_eq!(none, CompressionType::None);
    assert_eq!(lz4, CompressionType::LZ4);
    assert_eq!(zstd, CompressionType::Zstd);
    assert_eq!(gzip, CompressionType::Gzip);
}

/// Test: Compression type - None for no overhead
#[test]
fn test_compression_none() {
    let compression = CompressionType::None;

    // No compression for fastest I/O
    assert_eq!(compression, CompressionType::None);
}

/// Test: Compression type - LZ4 for speed
#[test]
fn test_compression_lz4_speed() {
    let compression = CompressionType::LZ4;

    // LZ4 is fastest compression
    assert_eq!(compression, CompressionType::LZ4);
}

/// Test: Compression type - Zstd for balance
#[test]
fn test_compression_zstd_balance() {
    let compression = CompressionType::Zstd;

    // Zstd balances speed and compression ratio
    assert_eq!(compression, CompressionType::Zstd);
}

/// Test: Compression type - Gzip for maximum compression
#[test]
fn test_compression_gzip_ratio() {
    let compression = CompressionType::Gzip;

    // Gzip for highest compression ratio
    assert_eq!(compression, CompressionType::Gzip);
}

/// Test: Compression type - serialization
#[test]
fn test_compression_type_serialization() {
    let compression = CompressionType::Zstd;

    // Serialize to JSON
    let json = serde_json::to_string(&compression).expect("Should serialize");

    // Deserialize back
    let deserialized: CompressionType =
        serde_json::from_str(&json).expect("Should deserialize");

    assert_eq!(compression, deserialized);
}

/// Test: Checkpoint metadata creation
#[test]
fn test_checkpoint_metadata_creation() {
    let metadata = CheckpointMetadata {
        checkpoint_id: "ckpt_001".to_string(),
        model_type: ModelType::DQN,
        model_version: "1.0.0".to_string(),
        created_at: chrono::Utc::now(),
        training_step: 1000,
        epoch: 10,
        learning_rate: 0.001,
        loss: 0.5,
        metrics: std::collections::HashMap::new(),
        hyperparameters: std::collections::HashMap::new(),
        compression: CompressionType::None,
        format: CheckpointFormat::Binary,
        file_size_bytes: 1024000,
        checksum: "abc123".to_string(),
    };

    // Verify basic fields
    assert_eq!(metadata.checkpoint_id, "ckpt_001");
    assert_eq!(metadata.model_type, ModelType::DQN);
    assert_eq!(metadata.model_version, "1.0.0");
    assert_eq!(metadata.training_step, 1000);
    assert_eq!(metadata.epoch, 10);
}

/// Test: Checkpoint metadata - training step validation
#[test]
fn test_checkpoint_metadata_training_step() {
    let metadata = CheckpointMetadata {
        checkpoint_id: "ckpt_002".to_string(),
        model_type: ModelType::MAMBA,
        model_version: "2.0.0".to_string(),
        created_at: chrono::Utc::now(),
        training_step: 5000,
        epoch: 50,
        learning_rate: 0.0001,
        loss: 0.3,
        metrics: std::collections::HashMap::new(),
        hyperparameters: std::collections::HashMap::new(),
        compression: CompressionType::LZ4,
        format: CheckpointFormat::Binary,
        file_size_bytes: 2048000,
        checksum: "def456".to_string(),
    };

    // Training step should be positive
    assert!(metadata.training_step > 0);
    assert!(metadata.epoch > 0);
}

/// Test: Checkpoint metadata - learning rate bounds
#[test]
fn test_checkpoint_metadata_learning_rate() {
    let metadata = CheckpointMetadata {
        checkpoint_id: "ckpt_003".to_string(),
        model_type: ModelType::TFT,
        model_version: "1.5.0".to_string(),
        created_at: chrono::Utc::now(),
        training_step: 10000,
        epoch: 100,
        learning_rate: 0.0005,
        loss: 0.2,
        metrics: std::collections::HashMap::new(),
        hyperparameters: std::collections::HashMap::new(),
        compression: CompressionType::Zstd,
        format: CheckpointFormat::JSON,
        file_size_bytes: 3072000,
        checksum: "ghi789".to_string(),
    };

    // Learning rate should be positive and reasonable
    assert!(metadata.learning_rate > 0.0);
    assert!(metadata.learning_rate <= 0.01);
}

/// Test: Checkpoint metadata - loss validation
#[test]
fn test_checkpoint_metadata_loss() {
    let metadata = CheckpointMetadata {
        checkpoint_id: "ckpt_004".to_string(),
        model_type: ModelType::TGNN,
        model_version: "1.2.0".to_string(),
        created_at: chrono::Utc::now(),
        training_step: 2000,
        epoch: 20,
        learning_rate: 0.001,
        loss: 0.15,
        metrics: std::collections::HashMap::new(),
        hyperparameters: std::collections::HashMap::new(),
        compression: CompressionType::None,
        format: CheckpointFormat::Binary,
        file_size_bytes: 1536000,
        checksum: "jkl012".to_string(),
    };

    // Loss should be non-negative
    assert!(metadata.loss >= 0.0);

    // Loss should be reasonable (not NaN or infinity)
    assert!(metadata.loss.is_finite());
}

/// Test: Checkpoint metadata - file size validation
#[test]
fn test_checkpoint_metadata_file_size() {
    let metadata = CheckpointMetadata {
        checkpoint_id: "ckpt_005".to_string(),
        model_type: ModelType::LiquidNN,
        model_version: "3.0.0".to_string(),
        created_at: chrono::Utc::now(),
        training_step: 15000,
        epoch: 150,
        learning_rate: 0.0002,
        loss: 0.1,
        metrics: std::collections::HashMap::new(),
        hyperparameters: std::collections::HashMap::new(),
        compression: CompressionType::Gzip,
        format: CheckpointFormat::MessagePack,
        file_size_bytes: 4096000,
        checksum: "mno345".to_string(),
    };

    // File size should be positive
    assert!(metadata.file_size_bytes > 0);

    // File size should be reasonable (not too large)
    assert!(metadata.file_size_bytes <= 10_000_000_000); // 10GB max
}

/// Test: Checkpoint metadata - checksum validation
#[test]
fn test_checkpoint_metadata_checksum() {
    let metadata = CheckpointMetadata {
        checkpoint_id: "ckpt_006".to_string(),
        model_type: ModelType::DQN,
        model_version: "1.1.0".to_string(),
        created_at: chrono::Utc::now(),
        training_step: 3000,
        epoch: 30,
        learning_rate: 0.0008,
        loss: 0.25,
        metrics: std::collections::HashMap::new(),
        hyperparameters: std::collections::HashMap::new(),
        compression: CompressionType::LZ4,
        format: CheckpointFormat::Binary,
        file_size_bytes: 2048000,
        checksum: "pqr678".to_string(),
    };

    // Checksum should not be empty
    assert!(!metadata.checksum.is_empty());

    // Checksum should be alphanumeric
    assert!(metadata.checksum.chars().all(|c| c.is_alphanumeric()));
}

/// Test: Checkpoint metadata - serialization roundtrip
#[test]
fn test_checkpoint_metadata_serialization() {
    let metadata = CheckpointMetadata {
        checkpoint_id: "ckpt_007".to_string(),
        model_type: ModelType::MAMBA,
        model_version: "2.1.0".to_string(),
        created_at: chrono::Utc::now(),
        training_step: 7000,
        epoch: 70,
        learning_rate: 0.0003,
        loss: 0.18,
        metrics: std::collections::HashMap::new(),
        hyperparameters: std::collections::HashMap::new(),
        compression: CompressionType::Zstd,
        format: CheckpointFormat::JSON,
        file_size_bytes: 3584000,
        checksum: "stu901".to_string(),
    };

    // Serialize to JSON
    let json = serde_json::to_string(&metadata).expect("Should serialize");

    // Deserialize back
    let deserialized: CheckpointMetadata =
        serde_json::from_str(&json).expect("Should deserialize");

    // Verify key fields match
    assert_eq!(metadata.checkpoint_id, deserialized.checkpoint_id);
    assert_eq!(metadata.model_type, deserialized.model_type);
    assert_eq!(metadata.model_version, deserialized.model_version);
    assert_eq!(metadata.training_step, deserialized.training_step);
    assert_eq!(metadata.epoch, deserialized.epoch);
}

/// Test: Model type variants
#[test]
fn test_model_type_variants() {
    // Test all model types
    let dqn = ModelType::DQN;
    let mamba = ModelType::MAMBA;
    let tft = ModelType::TFT;
    let tgnn = ModelType::TGNN;
    let liquid = ModelType::LiquidNN;

    // Verify equality
    assert_eq!(dqn, ModelType::DQN);
    assert_eq!(mamba, ModelType::MAMBA);
    assert_eq!(tft, ModelType::TFT);
    assert_eq!(tgnn, ModelType::TGNN);
    assert_eq!(liquid, ModelType::LiquidNN);
}

/// Test: Checkpoint metadata - metrics storage
#[test]
fn test_checkpoint_metadata_metrics() {
    let mut metrics = std::collections::HashMap::new();
    metrics.insert("accuracy".to_string(), 0.95);
    metrics.insert("precision".to_string(), 0.92);
    metrics.insert("recall".to_string(), 0.90);

    let metadata = CheckpointMetadata {
        checkpoint_id: "ckpt_008".to_string(),
        model_type: ModelType::TFT,
        model_version: "1.6.0".to_string(),
        created_at: chrono::Utc::now(),
        training_step: 12000,
        epoch: 120,
        learning_rate: 0.0004,
        loss: 0.12,
        metrics,
        hyperparameters: std::collections::HashMap::new(),
        compression: CompressionType::None,
        format: CheckpointFormat::Binary,
        file_size_bytes: 4608000,
        checksum: "vwx234".to_string(),
    };

    // Verify metrics are stored
    assert_eq!(metadata.metrics.len(), 3);
    assert_eq!(metadata.metrics.get("accuracy"), Some(&0.95));
    assert_eq!(metadata.metrics.get("precision"), Some(&0.92));
    assert_eq!(metadata.metrics.get("recall"), Some(&0.90));
}

/// Test: Checkpoint metadata - hyperparameters storage
#[test]
fn test_checkpoint_metadata_hyperparameters() {
    let mut hyperparameters = std::collections::HashMap::new();
    hyperparameters.insert("batch_size".to_string(), 32.0);
    hyperparameters.insert("dropout".to_string(), 0.1);
    hyperparameters.insert("num_layers".to_string(), 4.0);

    let metadata = CheckpointMetadata {
        checkpoint_id: "ckpt_009".to_string(),
        model_type: ModelType::TGNN,
        model_version: "1.3.0".to_string(),
        created_at: chrono::Utc::now(),
        training_step: 8000,
        epoch: 80,
        learning_rate: 0.0006,
        loss: 0.14,
        metrics: std::collections::HashMap::new(),
        hyperparameters,
        compression: CompressionType::LZ4,
        format: CheckpointFormat::JSON,
        file_size_bytes: 2560000,
        checksum: "yzA567".to_string(),
    };

    // Verify hyperparameters are stored
    assert_eq!(metadata.hyperparameters.len(), 3);
    assert_eq!(metadata.hyperparameters.get("batch_size"), Some(&32.0));
    assert_eq!(metadata.hyperparameters.get("dropout"), Some(&0.1));
    assert_eq!(metadata.hyperparameters.get("num_layers"), Some(&4.0));
}

/// Test: Checkpoint format - all formats compatible
#[test]
fn test_checkpoint_formats_compatibility() {
    let formats = vec![
        CheckpointFormat::Binary,
        CheckpointFormat::JSON,
        CheckpointFormat::MessagePack,
        CheckpointFormat::Custom,
    ];

    // All formats should be distinct
    for (i, format1) in formats.iter().enumerate() {
        for (j, format2) in formats.iter().enumerate() {
            if i == j {
                assert_eq!(format1, format2);
            } else {
                assert_ne!(format1, format2);
            }
        }
    }
}

/// Test: Compression types - all types compatible
#[test]
fn test_compression_types_compatibility() {
    let compressions = vec![
        CompressionType::None,
        CompressionType::LZ4,
        CompressionType::Zstd,
        CompressionType::Gzip,
    ];

    // All compression types should be distinct
    for (i, comp1) in compressions.iter().enumerate() {
        for (j, comp2) in compressions.iter().enumerate() {
            if i == j {
                assert_eq!(comp1, comp2);
            } else {
                assert_ne!(comp1, comp2);
            }
        }
    }
}