foxhunt/crates/ml-checkpoint/src/integration_tests.rs

//! Integration tests for the unified checkpoint system
//!
//! Comprehensive tests covering all functionality across all 5 AI models.

#[cfg(test)]
#[allow(
    clippy::assertions_on_result_states,
    clippy::doc_markdown,
    clippy::len_zero,
    clippy::str_to_string
)]
mod tests {
    use super::super::*;
    use crate::versioning::CompatibilityRisk;
    use std::sync::Arc;
    use tempfile::{tempdir, TempDir};

    // Production implementations for testing since we can't import the actual models
    // In a real implementation, these would be the actual model types

    #[derive(Debug)]
    struct MockModel {
        model_type: ModelType,
        name: String,
        version: String,
        state: Vec<u8>,
        hyperparams: HashMap<String, serde_json::Value>,
        metrics: HashMap<String, f64>,
    }

    impl MockModel {
        fn new(model_type: ModelType, name: &str, version: &str) -> Self {
            Self {
                model_type,
                name: name.to_string(),
                version: version.to_string(),
                state: vec![1, 2, 3, 4, 5],
                hyperparams: HashMap::new(),
                metrics: HashMap::new(),
            }
        }

        fn with_hyperparams(mut self, params: HashMap<String, serde_json::Value>) -> Self {
            self.hyperparams = params;
            self
        }

        fn with_metrics(mut self, metrics: HashMap<String, f64>) -> Self {
            self.metrics = metrics;
            self
        }
    }

    #[async_trait]
    impl Checkpointable for MockModel {
        fn model_type(&self) -> ModelType {
            self.model_type
        }

        fn model_name(&self) -> &str {
            &self.name
        }

        fn model_version(&self) -> &str {
            &self.version
        }

        async fn serialize_state(&self) -> Result<Vec<u8>, MLError> {
            Ok(self.state.clone())
        }

        async fn deserialize_state(&mut self, data: &[u8]) -> Result<(), MLError> {
            self.state = data.to_vec();
            Ok(())
        }

        fn get_training_state(&self) -> (Option<u64>, Option<u64>, Option<f64>, Option<f64>) {
            (Some(10), Some(1000), Some(0.1), Some(0.95))
        }

        fn get_hyperparameters(&self) -> HashMap<String, serde_json::Value> {
            self.hyperparams.clone()
        }

        fn get_metrics(&self) -> HashMap<String, f64> {
            self.metrics.clone()
        }

        fn get_architecture_info(&self) -> HashMap<String, serde_json::Value> {
            let mut info = HashMap::new();
            info.insert(
                "model_type".to_owned(),
                serde_json::Value::String(format!("{:?}", self.model_type)),
            );
            info.insert(
                "layers".to_owned(),
                serde_json::Value::Number(serde_json::Number::from(3)),
            );
            info
        }
    }

    /// Create a test checkpoint manager
    /// Returns (CheckpointManager, TempDir) to keep temp directory alive for test duration
    async fn create_test_manager() -> Result<(CheckpointManager, TempDir), MLError> {
        let temp_dir = tempdir().map_err(|e| {
            MLError::ModelError(format!("Failed to create temp directory in test: {}", e))
        })?;
        let config = CheckpointConfig {
            base_dir: temp_dir.path().to_path_buf(),
            compression: CompressionType::None,
            max_checkpoints_per_model: 3,
            auto_cleanup: false,
            ..Default::default()
        };

        let manager = CheckpointManager::new(config)?;
        Ok((manager, temp_dir))
    }

    #[tokio::test]
    async fn test_all_model_types_checkpoint() {
        let result = create_test_manager().await;
        assert!(
            result.is_ok(),
            "Failed to create test manager: {:?}",
            result.as_ref().err()
        );
        let (manager, _temp_dir) = result.unwrap();
        let model_types = [
            ModelType::DQN,
            ModelType::MAMBA,
            ModelType::TFT,
            ModelType::TGGN,
            ModelType::LNN,
        ];

        let mut checkpoint_ids = Vec::new();

        // Test saving checkpoints for all model types
        for (i, model_type) in model_types.into_iter().enumerate() {
            let mut model = MockModel::new(model_type, &format!("model_{}", i), "1.0.0");
            model.state = vec![i as u8; 10]; // Unique state for each model

            let checkpoint_result = manager
                .save_checkpoint(&model, Some(vec![format!("test_{}", i)]))
                .await;
            assert!(
                checkpoint_result.is_ok(),
                "Failed to save checkpoint: {:?}",
                checkpoint_result.err()
            );
            let checkpoint_id = checkpoint_result.unwrap();
            checkpoint_ids.push((model_type, checkpoint_id));
        }

        // Test loading checkpoints for all model types
        for (i, (model_type, checkpoint_id)) in checkpoint_ids.into_iter().enumerate() {
            let mut model = MockModel::new(model_type, &format!("model_{}", i), "1.0.0");
            let original_state = vec![i as u8; 10];

            // Change state before loading
            model.state = vec![99; 5];

            let load_result = manager.load_checkpoint(&mut model, &checkpoint_id).await;
            assert!(
                load_result.is_ok(),
                "Failed to load checkpoint: {:?}",
                load_result.err()
            );
            let metadata = load_result.unwrap();

            // Verify state was restored
            assert_eq!(model.state, original_state);
            assert_eq!(metadata.model_type, model_type);
            assert_eq!(metadata.model_name, format!("model_{}", i));
        }
    }

    #[tokio::test]
    async fn test_checkpoint_with_compression() {
        let temp_dir = tempdir().expect("Failed to create temp directory in test");
        let config = CheckpointConfig {
            base_dir: temp_dir.path().to_path_buf(),
            compression: CompressionType::Gzip,
            ..Default::default()
        };

        let manager =
            CheckpointManager::new(config).expect("Failed to create CheckpointManager in test");
        let mut model = MockModel::new(ModelType::DQN, "test_model", "1.0.0");

        // Create larger state for compression test
        model.state = vec![42; 1000];

        let checkpoint_id = manager
            .save_checkpoint(&model, None)
            .await
            .expect("Failed to save checkpoint in test");

        // Clear state
        model.state.clear();

        // Load and verify
        manager
            .load_checkpoint(&mut model, &checkpoint_id)
            .await
            .expect("Failed to get checkpoint info in test");
        assert_eq!(model.state, vec![42; 1000]);
    }

    #[tokio::test]
    async fn test_checkpoint_metadata_validation() {
        let (manager, _temp_dir) = create_test_manager()
            .await
            .map_err(|e| {
                panic!("Failed to create test manager: {}", e);
            })
            .unwrap();
        let mut hyperparams = HashMap::new();
        hyperparams.insert("learning_rate".to_owned(), serde_json::Value::from(0.001));
        hyperparams.insert("batch_size".to_owned(), serde_json::Value::from(32));

        let mut metrics = HashMap::new();
        metrics.insert("accuracy".to_owned(), 0.95);
        metrics.insert("loss".to_owned(), 0.05);

        let model = MockModel::new(ModelType::MAMBA, "test_model", "2.1.0")
            .with_hyperparams(hyperparams.clone())
            .with_metrics(metrics.clone());

        let _checkpoint_id = manager
            .save_checkpoint(&model, Some(vec!["validated".to_owned()]))
            .await
            .map_err(|e| {
                panic!("Operation failed in test: {}", e);
            })
            .unwrap();

        // Get checkpoint metadata
        let checkpoints = manager
            .list_checkpoints(ModelType::MAMBA, "test_model")
            .await;
        assert_eq!(checkpoints.len(), 1);

        let metadata = &checkpoints[0];
        assert_eq!(metadata.model_type, ModelType::MAMBA);
        assert_eq!(metadata.model_name, "test_model");
        assert_eq!(metadata.version, "2.1.0");
        assert_eq!(metadata.tags, vec!["validated".to_owned()]);
        assert!(metadata.hyperparameters.contains_key("learning_rate"));
        assert!(metadata.metrics.contains_key("accuracy"));
        assert_eq!(metadata.epoch, Some(10));
        assert_eq!(metadata.accuracy, Some(0.95));
    }

    #[tokio::test]
    async fn test_checkpoint_lifecycle_management() {
        let temp_dir = tempdir().expect("Failed to create temp directory in test");
        let config = CheckpointConfig {
            base_dir: temp_dir.path().to_path_buf(),
            max_checkpoints_per_model: 2,
            auto_cleanup: false, // Manual cleanup for testing
            ..Default::default()
        };

        let manager =
            CheckpointManager::new(config).expect("Failed to create CheckpointManager in test");
        let model = MockModel::new(ModelType::TFT, "lifecycle_test", "1.0.0");

        // Save multiple checkpoints
        let id1 = manager
            .save_checkpoint(&model, Some(vec!["v1".to_owned()]))
            .await
            .unwrap();
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;

        let id2 = manager
            .save_checkpoint(&model, Some(vec!["v2".to_owned()]))
            .await
            .unwrap();
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;

        let id3 = manager
            .save_checkpoint(&model, Some(vec!["v3".to_owned()]))
            .await
            .unwrap();

        // Should have 3 checkpoints before cleanup
        let checkpoints = manager
            .list_checkpoints(ModelType::TFT, "lifecycle_test")
            .await;
        assert_eq!(checkpoints.len(), 3);

        // Manual cleanup (simulating auto cleanup)
        // This would be done by the cleanup_old_checkpoints method

        // Test delete functionality
        manager
            .delete_checkpoint(&id1)
            .await
            .map_err(|e| {
                panic!("Failed to delete checkpoint: {}", e);
            })
            .unwrap();
        let checkpoints = manager
            .list_checkpoints(ModelType::TFT, "lifecycle_test")
            .await;
        assert_eq!(checkpoints.len(), 2);

        // Verify the correct checkpoint was deleted
        let remaining_ids: Vec<_> = checkpoints.iter().map(|c| &c.checkpoint_id).collect();
        assert!(remaining_ids.contains(&&id2));
        assert!(remaining_ids.contains(&&id3));
        assert!(!remaining_ids.contains(&&id1));
    }

    #[tokio::test]
    async fn test_checkpoint_search_and_filtering() {
        let (manager, _temp_dir) = create_test_manager().await.unwrap();

        // Create models with different tags
        let model1 = MockModel::new(ModelType::TGGN, "model_prod", "1.0.0");
        let model2 = MockModel::new(ModelType::TGGN, "model_dev", "1.1.0");
        let model3 = MockModel::new(ModelType::LNN, "model_test", "1.0.0");

        // Save with different tag combinations
        manager
            .save_checkpoint(
                &model1,
                Some(vec!["production".to_owned(), "validated".to_owned()]),
            )
            .await
            .unwrap();
        manager
            .save_checkpoint(&model2, Some(vec!["development".to_owned()]))
            .await
            .unwrap();
        manager
            .save_checkpoint(
                &model3,
                Some(vec!["test".to_owned(), "validated".to_owned()]),
            )
            .await
            .unwrap();

        // Test search by tags
        let production_checkpoints = manager
            .find_checkpoints_by_tags(&["production".to_owned()])
            .await;
        assert_eq!(production_checkpoints.len(), 1);
        assert_eq!(production_checkpoints[0].model_name, "model_prod");

        let validated_checkpoints = manager
            .find_checkpoints_by_tags(&["validated".to_owned()])
            .await;
        assert_eq!(validated_checkpoints.len(), 2);

        // Test list by model type
        let tggn_checkpoints = manager.list_checkpoints(ModelType::TGGN, "").await;
        assert_eq!(tggn_checkpoints.len(), 2);

        let lnn_checkpoints = manager.list_checkpoints(ModelType::LNN, "").await;
        assert_eq!(lnn_checkpoints.len(), 1);
    }

    #[tokio::test]
    async fn test_version_compatibility_checking() {
        let version_manager = VersionManager::new();

        // Test compatible versions
        let compat_info = version_manager
            .check_compatibility("1.0.0", "1.1.0", ModelType::DQN)
            .unwrap();

        assert!(compat_info.compatible);
        assert_eq!(compat_info.risk, CompatibilityRisk::Medium);

        // Test incompatible versions
        let incompat_info = version_manager
            .check_compatibility("1.0.0", "2.0.0", ModelType::DQN)
            .unwrap();

        assert!(!incompat_info.compatible);
        assert_eq!(incompat_info.risk, CompatibilityRisk::High);
        assert!(incompat_info.warnings.len() > 0);
    }

    #[tokio::test]
    async fn test_checkpoint_validation() {
        let (manager, _temp_dir) = create_test_manager().await.unwrap();
        let model = MockModel::new(ModelType::DQN, "validation_test", "1.0.0");

        let checkpoint_id = manager
            .save_checkpoint(&model, None)
            .await
            .map_err(|e| {
                panic!("Failed to save checkpoint: {}", e);
            })
            .unwrap();

        // Test normal loading (should pass validation)
        let mut model_copy = MockModel::new(ModelType::DQN, "validation_test", "1.0.0");
        let result = manager
            .load_checkpoint(&mut model_copy, &checkpoint_id)
            .await;
        assert!(result.is_ok());

        // Test loading with wrong model type (should fail)
        let mut wrong_model = MockModel::new(ModelType::MAMBA, "validation_test", "1.0.0");
        let result = manager
            .load_checkpoint(&mut wrong_model, &checkpoint_id)
            .await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_checkpoint_statistics() {
        let (manager, _temp_dir) = create_test_manager().await.unwrap();

        // Initial stats should be zero
        let initial_stats = manager.get_stats();
        assert_eq!(initial_stats.get("total_saved").unwrap_or(&0), &0);
        assert_eq!(initial_stats.get("total_loaded").unwrap_or(&0), &0);

        // Save a checkpoint
        let model = MockModel::new(ModelType::TFT, "stats_test", "1.0.0");
        let checkpoint_id = manager.save_checkpoint(&model, None).await.unwrap();

        // Stats should reflect the save
        let save_stats = manager.get_stats();
        assert_eq!(save_stats.get("total_saved").unwrap_or(&0), &1);
        assert!(save_stats.get("total_bytes_saved").unwrap_or(&0) > &0);

        // Load the checkpoint
        let mut model_copy = MockModel::new(ModelType::TFT, "stats_test", "1.0.0");
        manager
            .load_checkpoint(&mut model_copy, &checkpoint_id)
            .await
            .unwrap();

        // Stats should reflect both save and load
        let final_stats = manager.get_stats();
        assert_eq!(final_stats.get("total_saved").unwrap_or(&0), &1);
        assert_eq!(final_stats.get("total_loaded").unwrap_or(&0), &1);
        assert!(final_stats.get("total_bytes_loaded").unwrap_or(&0) > &0);
    }

    #[tokio::test]
    async fn test_concurrent_checkpoint_operations() {
        let (manager_inner, _temp_dir) = create_test_manager()
            .await
            .map_err(|e| {
                panic!("Failed to create test manager: {}", e);
            })
            .unwrap();
        let manager: Arc<CheckpointManager> = Arc::new(manager_inner);
        let mut handles = Vec::new();

        // Start multiple concurrent save operations
        for i in 0..5 {
            let manager_clone = Arc::clone(&manager);
            let handle = tokio::spawn(async move {
                let model =
                    MockModel::new(ModelType::DQN, &format!("concurrent_model_{}", i), "1.0.0");
                manager_clone
                    .save_checkpoint(&model, Some(vec![format!("concurrent_{}", i)]))
                    .await
            });
            handles.push(handle);
        }

        // Wait for all operations to complete
        let mut checkpoint_ids = Vec::new();
        for handle in handles {
            let checkpoint_id = handle
                .await
                .map_err(|e| {
                    panic!("Join handle failed: {}", e);
                })
                .unwrap()
                .map_err(|e| {
                    panic!("Save checkpoint failed: {}", e);
                })
                .unwrap();
            checkpoint_ids.push(checkpoint_id);
        }

        // Verify all checkpoints were saved
        assert_eq!(checkpoint_ids.len(), 5);

        // Test concurrent loads
        let mut load_handles = Vec::new();
        for (i, checkpoint_id) in checkpoint_ids.into_iter().enumerate() {
            let manager_clone: Arc<CheckpointManager> = Arc::clone(&manager);
            let handle = tokio::spawn(async move {
                let mut model =
                    MockModel::new(ModelType::DQN, &format!("concurrent_model_{}", i), "1.0.0");
                manager_clone
                    .load_checkpoint(&mut model, &checkpoint_id)
                    .await
            });
            load_handles.push(handle);
        }

        // Verify all loads succeed
        for handle in load_handles {
            assert!(handle
                .await
                .map_err(|e| {
                    panic!("Join handle failed: {}", e);
                })
                .unwrap()
                .is_ok());
        }
    }

    #[tokio::test]
    async fn test_latest_checkpoint_functionality() {
        let (manager, _temp_dir) = create_test_manager().await.unwrap();
        let model = MockModel::new(ModelType::MAMBA, "latest_test", "1.0.0");

        // Initially no latest checkpoint
        let mut model_copy = MockModel::new(ModelType::MAMBA, "latest_test", "1.0.0");
        let latest = manager
            .load_latest_checkpoint(&mut model_copy)
            .await
            .unwrap();
        assert!(latest.is_none());

        // Save first checkpoint
        let _id1 = manager
            .save_checkpoint(&model, Some(vec!["first".to_owned()]))
            .await
            .unwrap();
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;

        // Save second checkpoint (should become latest)
        let id2 = manager
            .save_checkpoint(&model, Some(vec!["second".to_owned()]))
            .await
            .unwrap();

        // Test latest checkpoint loading
        let mut model_copy = MockModel::new(ModelType::MAMBA, "latest_test", "1.0.0");
        let latest = manager
            .load_latest_checkpoint(&mut model_copy)
            .await
            .unwrap();

        assert!(latest.is_some());
        let latest_metadata = latest.unwrap();
        assert_eq!(latest_metadata.checkpoint_id, id2);
        assert!(latest_metadata.tags.contains(&"second".to_owned()));
    }
}