foxhunt/services/ml_training_service/tests/checkpoint_manager_tests.rs

//! # TDD Checkpoint Manager Tests
//!
//! Test-Driven Development (TDD) tests for checkpoint retention, versioning, and cleanup.
//! All tests should FAIL initially, then pass after implementation.
//!
//! ## Test Coverage
//! - Retention policy (keep best 5 checkpoints per model)
//! - Automatic cleanup (>30 days old)
//! - Semantic versioning (v1.0.0, v1.0.1)
//! - SHA256 integrity validation
//! - Database integration (ml_model_versions table)

use chrono::{Duration, Utc};
use ml::checkpoint::{CheckpointMetadata, CheckpointFormat, CompressionType};
use ml::ModelType;
use ml_training_service::checkpoint_manager::{CheckpointManager, RetentionPolicy};
use sqlx::PgPool;
use std::collections::HashMap;

/// Helper to create test checkpoint metadata
fn create_test_metadata(
    model_type: ModelType,
    model_name: &str,
    version: &str,
    accuracy: f64,
    sharpe_ratio: f64,
    created_days_ago: i64,
) -> CheckpointMetadata {
    let mut metrics = HashMap::new();
    metrics.insert("accuracy".to_string(), accuracy);
    metrics.insert("sharpe_ratio".to_string(), sharpe_ratio);

    CheckpointMetadata {
        checkpoint_id: uuid::Uuid::new_v4().to_string(),
        model_type,
        model_name: model_name.to_string(),
        version: version.to_string(),
        created_at: Utc::now() - Duration::days(created_days_ago),
        epoch: Some(100),
        step: Some(10000),
        loss: Some(0.05),
        accuracy: Some(accuracy),
        hyperparameters: HashMap::new(),
        metrics,
        architecture: HashMap::new(),
        format: CheckpointFormat::Binary,
        compression: CompressionType::LZ4,
        file_size: 1024 * 1024, // 1MB
        compressed_size: Some(512 * 1024), // 512KB
        checksum: "0".repeat(64), // Placeholder SHA256
        tags: vec![],
        custom_metadata: HashMap::new(),
        signature: None,
        signature_algorithm: "none".to_string(),
        signing_key_id: "test".to_string(),
        signed_at: None,
    }
}

/// Helper to setup test database
async fn setup_test_db() -> PgPool {
    let database_url = std::env::var("DATABASE_URL")
        .unwrap_or_else(|_| "postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt".to_string());

    PgPool::connect(&database_url)
        .await
        .expect("Failed to connect to test database")
}

/// Helper to cleanup test data
async fn cleanup_test_data(pool: &PgPool, model_name: &str) {
    let _ = sqlx::query!(
        "DELETE FROM ml_model_versions WHERE metadata->>'test_model_name' = $1",
        model_name
    )
    .execute(pool)
    .await;
}

// ================================================================================================
// TEST 1: Retention Policy - Keep Best 5 Checkpoints
// ================================================================================================

#[tokio::test]
async fn test_retention_policy_keeps_best_5_checkpoints() {
    let pool = setup_test_db().await;
    let test_model_name = "test_retention_dqn";

    // Cleanup before test
    cleanup_test_data(&pool, test_model_name).await;

    let retention_policy = RetentionPolicy {
        max_checkpoints_per_model: 5,
        ranking_metric: "sharpe_ratio".to_string(),
        ascending: false, // Higher is better
    };

    let manager = CheckpointManager::new(pool.clone(), retention_policy)
        .await
        .expect("Failed to create CheckpointManager");

    // Create 10 checkpoints with different Sharpe ratios
    let sharpe_ratios = vec![1.2, 2.5, 1.8, 3.0, 1.5, 2.2, 1.9, 2.8, 1.7, 2.1];

    for (i, sharpe) in sharpe_ratios.iter().enumerate() {
        let metadata = create_test_metadata(
            ModelType::DQN,
            test_model_name,
            &format!("1.0.{}", i),
            0.85,
            *sharpe,
            0, // All created today
        );

        manager
            .register_checkpoint(metadata)
            .await
            .expect("Failed to register checkpoint");
    }

    // Apply retention policy
    manager
        .apply_retention_policy(ModelType::DQN, test_model_name)
        .await
        .expect("Failed to apply retention policy");

    // Verify only 5 best checkpoints remain
    let remaining = manager
        .list_checkpoints(ModelType::DQN, test_model_name)
        .await
        .expect("Failed to list checkpoints");

    assert_eq!(remaining.len(), 5, "Should keep exactly 5 checkpoints");

    // Verify they are the top 5 by Sharpe ratio
    let mut remaining_sharpe: Vec<f64> = remaining
        .iter()
        .map(|m| m.metrics.get("sharpe_ratio").copied().unwrap_or(0.0))
        .collect();
    remaining_sharpe.sort_by(|a, b| b.partial_cmp(a).unwrap());

    let expected_top_5 = vec![3.0, 2.8, 2.5, 2.2, 2.1];
    assert_eq!(
        remaining_sharpe, expected_top_5,
        "Should keep checkpoints with highest Sharpe ratios"
    );

    // Cleanup
    cleanup_test_data(&pool, test_model_name).await;
}

// ================================================================================================
// TEST 2: Automatic Cleanup - Remove Checkpoints Older Than 30 Days
// ================================================================================================

#[tokio::test]
async fn test_automatic_cleanup_old_checkpoints() {
    let pool = setup_test_db().await;
    let test_model_name = "test_cleanup_mamba";

    cleanup_test_data(&pool, test_model_name).await;

    let retention_policy = RetentionPolicy {
        max_checkpoints_per_model: 10,
        ranking_metric: "accuracy".to_string(),
        ascending: false,
    };

    let manager = CheckpointManager::new(pool.clone(), retention_policy)
        .await
        .expect("Failed to create CheckpointManager");

    // Create checkpoints with different ages
    let test_data = vec![
        (5, 0.90),   // 5 days old
        (15, 0.88),  // 15 days old
        (25, 0.92),  // 25 days old
        (35, 0.89),  // 35 days old - should be removed
        (40, 0.91),  // 40 days old - should be removed
        (50, 0.87),  // 50 days old - should be removed
    ];

    for (i, (days_ago, accuracy)) in test_data.iter().enumerate() {
        let metadata = create_test_metadata(
            ModelType::MAMBA,
            test_model_name,
            &format!("1.0.{}", i),
            *accuracy,
            1.5,
            *days_ago,
        );

        manager
            .register_checkpoint(metadata)
            .await
            .expect("Failed to register checkpoint");
    }

    // Apply cleanup (30-day threshold)
    let cleanup_count = manager
        .cleanup_old_checkpoints(ModelType::MAMBA, test_model_name, 30)
        .await
        .expect("Failed to cleanup old checkpoints");

    assert_eq!(cleanup_count, 3, "Should remove 3 checkpoints older than 30 days");

    // Verify only recent checkpoints remain
    let remaining = manager
        .list_checkpoints(ModelType::MAMBA, test_model_name)
        .await
        .expect("Failed to list checkpoints");

    assert_eq!(remaining.len(), 3, "Should have 3 remaining checkpoints");

    for checkpoint in &remaining {
        let age_days = (Utc::now() - checkpoint.created_at).num_days();
        assert!(
            age_days < 30,
            "Checkpoint should be less than 30 days old, got {} days",
            age_days
        );
    }

    cleanup_test_data(&pool, test_model_name).await;
}

// ================================================================================================
// TEST 3: Semantic Versioning Validation
// ================================================================================================

#[tokio::test]
async fn test_semantic_versioning() {
    let pool = setup_test_db().await;
    let test_model_name = "test_versioning_ppo";

    cleanup_test_data(&pool, test_model_name).await;

    let retention_policy = RetentionPolicy::default();
    let manager = CheckpointManager::new(pool.clone(), retention_policy)
        .await
        .expect("Failed to create CheckpointManager");

    // Test valid semantic versions
    let valid_versions = vec!["1.0.0", "1.0.1", "1.1.0", "2.0.0", "1.0.0-alpha", "1.0.0-beta+build1"];

    for (i, version) in valid_versions.iter().enumerate() {
        let metadata = create_test_metadata(
            ModelType::PPO,
            test_model_name,
            version,
            0.85,
            1.5,
            0,
        );

        let result = manager.register_checkpoint(metadata).await;
        assert!(
            result.is_ok(),
            "Valid semantic version '{}' should be accepted",
            version
        );
    }

    // Test invalid semantic versions
    let invalid_versions = vec!["1.0", "v1.0.0", "1.0.0.0", "1.a.0", ""];

    for version in &invalid_versions {
        let metadata = create_test_metadata(
            ModelType::PPO,
            &format!("{}_invalid", test_model_name),
            version,
            0.85,
            1.5,
            0,
        );

        let result = manager.validate_version(&metadata.version).await;
        assert!(
            result.is_err(),
            "Invalid semantic version '{}' should be rejected",
            version
        );
    }

    cleanup_test_data(&pool, test_model_name).await;
}

// ================================================================================================
// TEST 4: SHA256 Integrity Validation
// ================================================================================================

#[tokio::test]
async fn test_sha256_integrity_validation() {
    let pool = setup_test_db().await;
    let test_model_name = "test_integrity_tft";

    cleanup_test_data(&pool, test_model_name).await;

    let retention_policy = RetentionPolicy::default();
    let manager = CheckpointManager::new(pool.clone(), retention_policy)
        .await
        .expect("Failed to create CheckpointManager");

    // Create checkpoint with known data
    let test_data = b"test checkpoint data for integrity validation";
    let expected_checksum = format!("{:x}", sha2::Sha256::digest(test_data));

    let mut metadata = create_test_metadata(
        ModelType::TFT,
        test_model_name,
        "1.0.0",
        0.90,
        2.0,
        0,
    );
    metadata.checksum = expected_checksum.clone();

    // Register checkpoint
    manager
        .register_checkpoint(metadata.clone())
        .await
        .expect("Failed to register checkpoint");

    // Validate integrity with correct data
    let valid_result = manager
        .validate_checksum(&metadata.checkpoint_id, test_data)
        .await;
    assert!(valid_result.is_ok(), "Valid checksum should pass");

    // Validate integrity with corrupted data
    let corrupted_data = b"corrupted data";
    let invalid_result = manager
        .validate_checksum(&metadata.checkpoint_id, corrupted_data)
        .await;
    assert!(
        invalid_result.is_err(),
        "Invalid checksum should fail validation"
    );

    cleanup_test_data(&pool, test_model_name).await;
}

// ================================================================================================
// TEST 5: Database Integration - ml_model_versions Table
// ================================================================================================

#[tokio::test]
async fn test_database_integration() {
    let pool = setup_test_db().await;
    let test_model_name = "test_db_integration_dqn";

    cleanup_test_data(&pool, test_model_name).await;

    let retention_policy = RetentionPolicy::default();
    let manager = CheckpointManager::new(pool.clone(), retention_policy)
        .await
        .expect("Failed to create CheckpointManager");

    // Create and register checkpoint
    let metadata = create_test_metadata(
        ModelType::DQN,
        test_model_name,
        "1.0.0",
        0.95,
        2.5,
        0,
    );

    let checkpoint_id = manager
        .register_checkpoint(metadata.clone())
        .await
        .expect("Failed to register checkpoint");

    // Verify checkpoint was inserted into database
    let result = sqlx::query!(
        r#"
        SELECT model_id, model_type, version, checksum, metrics
        FROM ml_model_versions
        WHERE model_id = $1
        "#,
        checkpoint_id
    )
    .fetch_one(&pool)
    .await;

    assert!(result.is_ok(), "Checkpoint should exist in database");

    let record = result.unwrap();
    assert_eq!(record.model_type, "DQN");
    assert_eq!(record.version, "1.0.0");
    assert_eq!(record.checksum, metadata.checksum);

    // Verify metrics are stored correctly
    let metrics: serde_json::Value = record.metrics;
    assert_eq!(
        metrics["accuracy"].as_f64().unwrap(),
        0.95,
        "Accuracy should match"
    );
    assert_eq!(
        metrics["sharpe_ratio"].as_f64().unwrap(),
        2.5,
        "Sharpe ratio should match"
    );

    cleanup_test_data(&pool, test_model_name).await;
}

// ================================================================================================
// TEST 6: Combined Retention + Cleanup Workflow
// ================================================================================================

#[tokio::test]
async fn test_combined_retention_and_cleanup() {
    let pool = setup_test_db().await;
    let test_model_name = "test_combined_workflow";

    cleanup_test_data(&pool, test_model_name).await;

    let retention_policy = RetentionPolicy {
        max_checkpoints_per_model: 3,
        ranking_metric: "sharpe_ratio".to_string(),
        ascending: false,
    };

    let manager = CheckpointManager::new(pool.clone(), retention_policy)
        .await
        .expect("Failed to create CheckpointManager");

    // Create 8 checkpoints with different ages and Sharpe ratios
    let test_data = vec![
        (5, 2.5),   // Recent, high Sharpe - KEEP
        (10, 3.0),  // Recent, highest Sharpe - KEEP
        (15, 2.2),  // Recent, medium Sharpe - KEEP
        (20, 1.8),  // Recent, low Sharpe - REMOVE (retention)
        (35, 2.8),  // Old, high Sharpe - REMOVE (30-day rule)
        (40, 1.5),  // Old, low Sharpe - REMOVE (30-day rule)
        (45, 2.0),  // Old, medium Sharpe - REMOVE (30-day rule)
        (50, 3.2),  // Old, highest Sharpe - REMOVE (30-day rule)
    ];

    for (i, (days_ago, sharpe)) in test_data.iter().enumerate() {
        let metadata = create_test_metadata(
            ModelType::DQN,
            test_model_name,
            &format!("1.0.{}", i),
            0.85,
            *sharpe,
            *days_ago,
        );

        manager
            .register_checkpoint(metadata)
            .await
            .expect("Failed to register checkpoint");
    }

    // Apply 30-day cleanup first
    let cleanup_count = manager
        .cleanup_old_checkpoints(ModelType::DQN, test_model_name, 30)
        .await
        .expect("Failed to cleanup old checkpoints");

    assert_eq!(
        cleanup_count, 4,
        "Should remove 4 checkpoints older than 30 days"
    );

    // Then apply retention policy (keep best 3)
    manager
        .apply_retention_policy(ModelType::DQN, test_model_name)
        .await
        .expect("Failed to apply retention policy");

    // Verify final state: 3 checkpoints, all recent, highest Sharpe ratios
    let remaining = manager
        .list_checkpoints(ModelType::DQN, test_model_name)
        .await
        .expect("Failed to list checkpoints");

    assert_eq!(remaining.len(), 3, "Should have exactly 3 checkpoints remaining");

    // Verify all are recent (<30 days)
    for checkpoint in &remaining {
        let age_days = (Utc::now() - checkpoint.created_at).num_days();
        assert!(age_days < 30, "All remaining checkpoints should be recent");
    }

    // Verify they have the top 3 Sharpe ratios among recent checkpoints
    let mut sharpe_ratios: Vec<f64> = remaining
        .iter()
        .map(|m| m.metrics.get("sharpe_ratio").copied().unwrap_or(0.0))
        .collect();
    sharpe_ratios.sort_by(|a, b| b.partial_cmp(a).unwrap());

    let expected = vec![3.0, 2.5, 2.2];
    assert_eq!(sharpe_ratios, expected, "Should keep top 3 recent checkpoints");

    cleanup_test_data(&pool, test_model_name).await;
}

// ================================================================================================
// TEST 7: Version Comparison and Latest Checkpoint
// ================================================================================================

#[tokio::test]
async fn test_version_comparison() {
    let pool = setup_test_db().await;
    let test_model_name = "test_version_compare";

    cleanup_test_data(&pool, test_model_name).await;

    let retention_policy = RetentionPolicy::default();
    let manager = CheckpointManager::new(pool.clone(), retention_policy)
        .await
        .expect("Failed to create CheckpointManager");

    // Create checkpoints with different versions
    let versions = vec!["1.0.0", "1.0.1", "1.1.0", "2.0.0"];

    for (i, version) in versions.iter().enumerate() {
        let metadata = create_test_metadata(
            ModelType::DQN,
            test_model_name,
            version,
            0.85,
            1.5,
            0,
        );

        manager
            .register_checkpoint(metadata)
            .await
            .expect("Failed to register checkpoint");
    }

    // Get latest version
    let latest = manager
        .get_latest_checkpoint(ModelType::DQN, test_model_name)
        .await
        .expect("Failed to get latest checkpoint");

    assert!(latest.is_some(), "Should find latest checkpoint");
    assert_eq!(
        latest.unwrap().version,
        "2.0.0",
        "Latest version should be 2.0.0"
    );

    cleanup_test_data(&pool, test_model_name).await;
}