#![allow( unused_variables, clippy::unwrap_used, clippy::expect_used, clippy::indexing_slicing, clippy::useless_vec )] //! # 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::{CheckpointFormat, CheckpointMetadata, CompressionType}; use ml::ModelType; use ml_training_service::checkpoint_manager::{CheckpointManager, RetentionPolicy}; use sha2::{Digest, Sha256}; 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", ) .bind(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 = 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_or(std::cmp::Ordering::Equal)); 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 mut hasher = Sha256::new(); hasher.update(test_data); let expected_checksum = format!("{:x}", hasher.finalize()); 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: Result<(uuid::Uuid, String, String, String, serde_json::Value), _> = sqlx::query_as( r#" SELECT model_id, model_type, version, checksum, metrics FROM ml_model_versions WHERE model_id = $1 "#, ) .bind(checkpoint_id) .fetch_one(&pool) .await; assert!(result.is_ok(), "Checkpoint should exist in database"); let record = result.unwrap(); assert_eq!(record.1, "DQN"); assert_eq!(record.2, "1.0.0"); assert_eq!(record.3, metadata.checksum); // Verify metrics are stored correctly let metrics: serde_json::Value = record.4; 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 = 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_or(std::cmp::Ordering::Equal)); 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; }