foxhunt/storage/tests/checkpoint_archival_tests.rs

//! Checkpoint archival and backup/restore operation tests
//!
//! Tests comprehensive checkpoint management scenarios including:
//! - Large checkpoint uploads
//! - Checkpoint retrieval with verification
//! - Backup and restore workflows
//! - Concurrent checkpoint operations
//! - Error handling for checkpoint operations

use object_store::memory::InMemory;
use object_store::ObjectStore;
use std::sync::Arc;
use storage::object_store_backend::ObjectStoreBackend;
use storage::Storage;

/// Helper to create test backend with in-memory store
fn create_test_backend() -> ObjectStoreBackend {
    let in_memory_store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
    storage::object_store_backend::test_helpers::new_for_testing(
        in_memory_store,
        "checkpoints-bucket".to_string(),
    )
}

#[tokio::test]
async fn test_checkpoint_upload_and_download() {
    let backend = create_test_backend();

    // Simulate checkpoint data (10MB)
    let checkpoint_data = vec![0xAB; 10 * 1024 * 1024];
    let checkpoint_path = backend.get_checkpoint_path("mamba2", "epoch_100");

    // Upload checkpoint
    backend
        .store(&checkpoint_path, &checkpoint_data)
        .await
        .expect("Failed to upload checkpoint");

    // Download checkpoint
    let downloaded = backend
        .retrieve(&checkpoint_path)
        .await
        .expect("Failed to download checkpoint");

    assert_eq!(
        downloaded.len(),
        checkpoint_data.len(),
        "Downloaded checkpoint size mismatch"
    );
    assert_eq!(
        downloaded, checkpoint_data,
        "Downloaded checkpoint data mismatch"
    );
}

#[tokio::test]
async fn test_checkpoint_metadata_storage() {
    let backend = create_test_backend();

    let checkpoint_path = backend.get_checkpoint_path("dqn", "epoch_50");
    let metadata_path = backend.get_metadata_path("dqn", "v1.0");

    // Store checkpoint
    let checkpoint_data = b"checkpoint weights";
    backend
        .store(&checkpoint_path, checkpoint_data)
        .await
        .unwrap();

    // Store metadata
    let metadata = serde_json::json!({
        "model": "dqn",
        "epoch": 50,
        "loss": 0.123,
        "accuracy": 0.95
    });
    let metadata_bytes = serde_json::to_vec(&metadata).unwrap();
    backend
        .store(&metadata_path, &metadata_bytes)
        .await
        .unwrap();

    // Verify both exist
    assert!(backend.exists(&checkpoint_path).await.unwrap());
    assert!(backend.exists(&metadata_path).await.unwrap());

    // Retrieve and verify metadata
    let retrieved_metadata = backend.retrieve(&metadata_path).await.unwrap();
    let parsed: serde_json::Value = serde_json::from_slice(&retrieved_metadata).unwrap();
    assert_eq!(parsed["model"], "dqn");
    assert_eq!(parsed["epoch"], 50);
}

#[tokio::test]
async fn test_checkpoint_backup_workflow() {
    let backend = create_test_backend();

    // Primary checkpoint
    let primary_path = "models/ppo/checkpoints/primary/epoch_100.safetensors";
    let checkpoint_data = vec![0xCD; 5 * 1024 * 1024]; // 5MB

    // Store primary checkpoint
    backend.store(primary_path, &checkpoint_data).await.unwrap();

    // Create backup
    let backup_path = "models/ppo/backups/epoch_100_backup.safetensors";
    let retrieved = backend.retrieve(primary_path).await.unwrap();
    backend.store(backup_path, &retrieved).await.unwrap();

    // Verify both exist
    assert!(backend.exists(primary_path).await.unwrap());
    assert!(backend.exists(backup_path).await.unwrap());

    // Verify backup integrity
    let backup_data = backend.retrieve(backup_path).await.unwrap();
    assert_eq!(backup_data, checkpoint_data);
}

#[tokio::test]
async fn test_checkpoint_restore_from_backup() {
    let backend = create_test_backend();

    let backup_path = "backups/tft/epoch_200.safetensors";
    let restore_path = "models/tft/checkpoints/epoch_200_restored.safetensors";

    let backup_data = vec![0xEF; 8 * 1024 * 1024]; // 8MB

    // Store backup
    backend.store(backup_path, &backup_data).await.unwrap();

    // Simulate restore operation
    let restored_data = backend.retrieve(backup_path).await.unwrap();
    backend.store(restore_path, &restored_data).await.unwrap();

    // Verify restored checkpoint
    let final_data = backend.retrieve(restore_path).await.unwrap();
    assert_eq!(final_data.len(), backup_data.len());
    assert_eq!(final_data, backup_data);
}

#[tokio::test]
async fn test_checkpoint_versioning() {
    let backend = create_test_backend();

    let model_name = "mamba2";
    let versions = ["v1.0", "v1.1", "v2.0"];

    // Store multiple checkpoint versions
    for version in &versions {
        let path = backend.get_checkpoint_path(model_name, version);
        let data = format!("checkpoint_{}", version).into_bytes();
        backend.store(&path, &data).await.unwrap();
    }

    // List all checkpoints
    let prefix = format!("models/{}/checkpoints/", model_name);
    let checkpoints = backend.list(&prefix).await.unwrap();

    assert_eq!(checkpoints.len(), versions.len());

    // Verify each version exists
    for version in &versions {
        let path = backend.get_checkpoint_path(model_name, version);
        assert!(backend.exists(&path).await.unwrap());
    }
}

#[tokio::test]
async fn test_checkpoint_deletion() {
    let backend = create_test_backend();

    let checkpoint_path = backend.get_checkpoint_path("dqn", "old_epoch_10");
    let checkpoint_data = b"old checkpoint data";

    // Store checkpoint
    backend
        .store(&checkpoint_path, checkpoint_data)
        .await
        .unwrap();
    assert!(backend.exists(&checkpoint_path).await.unwrap());

    // Delete checkpoint
    let deleted = backend.delete(&checkpoint_path).await.unwrap();
    assert!(deleted);

    // Verify deletion
    assert!(!backend.exists(&checkpoint_path).await.unwrap());
}

#[tokio::test]
async fn test_checkpoint_cleanup_old_versions() {
    let backend = create_test_backend();

    let model_name = "ppo";
    let max_checkpoints = 3;

    // Store 5 checkpoints (should keep only 3 latest)
    for i in 1..=5 {
        let path = backend.get_checkpoint_path(model_name, &format!("epoch_{}", i * 10));
        let data = format!("checkpoint_{}", i).into_bytes();
        backend.store(&path, &data).await.unwrap();
    }

    let prefix = format!("models/{}/checkpoints/", model_name);
    let all_checkpoints = backend.list(&prefix).await.unwrap();
    assert_eq!(all_checkpoints.len(), 5);

    // Simulate cleanup: delete oldest checkpoints
    let mut to_delete = all_checkpoints.clone();
    to_delete.sort();
    while to_delete.len() > max_checkpoints {
        let oldest = to_delete.remove(0);
        backend.delete(&oldest).await.unwrap();
    }

    // Verify only max_checkpoints remain
    let remaining = backend.list(&prefix).await.unwrap();
    assert_eq!(remaining.len(), max_checkpoints);
}

#[tokio::test]
async fn test_concurrent_checkpoint_operations() {
    let backend = Arc::new(create_test_backend());
    let mut handles = vec![];

    // Concurrent checkpoint uploads
    for i in 0..5 {
        let backend = Arc::clone(&backend);
        let handle = tokio::spawn(async move {
            let path = format!("models/concurrent_test/checkpoints/epoch_{}.bin", i);
            let data = vec![i as u8; 1024 * 1024]; // 1MB each
            backend.store(&path, &data).await.unwrap();

            // Verify upload
            let retrieved = backend.retrieve(&path).await.unwrap();
            assert_eq!(retrieved.len(), data.len());
        });
        handles.push(handle);
    }

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

    // Verify all checkpoints exist
    let checkpoints = backend
        .list("models/concurrent_test/checkpoints/")
        .await
        .unwrap();
    assert_eq!(checkpoints.len(), 5);
}

#[tokio::test]
async fn test_checkpoint_integrity_verification() {
    let backend = create_test_backend();

    let checkpoint_path = backend.get_checkpoint_path("tft", "epoch_150");
    let checkpoint_data = vec![0x42; 15 * 1024 * 1024]; // 15MB

    // Calculate expected checksum (SHA-256)
    use sha2::{Digest, Sha256};
    let mut hasher = Sha256::new();
    hasher.update(&checkpoint_data);
    let expected_checksum = format!("{:x}", hasher.finalize());

    // Store checkpoint
    backend
        .store(&checkpoint_path, &checkpoint_data)
        .await
        .unwrap();

    // Retrieve and verify checksum
    let retrieved = backend.retrieve(&checkpoint_path).await.unwrap();
    let mut hasher = Sha256::new();
    hasher.update(&retrieved);
    let actual_checksum = format!("{:x}", hasher.finalize());

    assert_eq!(actual_checksum, expected_checksum);
}

#[tokio::test]
async fn test_checkpoint_partial_upload_failure() {
    let backend = create_test_backend();

    let checkpoint_path = backend.get_checkpoint_path("mamba2", "epoch_partial");
    let checkpoint_data = vec![0xAA; 20 * 1024 * 1024]; // 20MB

    // Upload checkpoint
    let result = backend.store(&checkpoint_path, &checkpoint_data).await;
    assert!(result.is_ok());

    // Verify metadata reflects correct size
    let metadata = backend.metadata(&checkpoint_path).await.unwrap();
    assert_eq!(metadata.size, checkpoint_data.len() as u64);
}

#[tokio::test]
async fn test_checkpoint_list_with_pagination() {
    let backend = create_test_backend();

    // Store 20 checkpoints
    for i in 1..=20 {
        let path = format!("models/pagination_test/checkpoints/epoch_{}.bin", i);
        backend.store(&path, &vec![i as u8; 1024]).await.unwrap();
    }

    // List all checkpoints
    let all_checkpoints = backend
        .list("models/pagination_test/checkpoints/")
        .await
        .unwrap();

    assert_eq!(all_checkpoints.len(), 20);
}

#[tokio::test]
async fn test_checkpoint_empty_content() {
    let backend = create_test_backend();

    let checkpoint_path = backend.get_checkpoint_path("empty_test", "epoch_0");
    let empty_data = b"";

    // Store empty checkpoint
    backend.store(&checkpoint_path, empty_data).await.unwrap();

    // Retrieve and verify
    let retrieved = backend.retrieve(&checkpoint_path).await.unwrap();
    assert_eq!(retrieved.len(), 0);

    // Verify metadata
    let metadata = backend.metadata(&checkpoint_path).await.unwrap();
    assert_eq!(metadata.size, 0);
}

#[tokio::test]
async fn test_checkpoint_overwrite_protection() {
    let backend = create_test_backend();

    let checkpoint_path = backend.get_checkpoint_path("overwrite_test", "epoch_100");
    let original_data = b"original checkpoint v1";
    let new_data = b"updated checkpoint v2";

    // Store original
    backend
        .store(&checkpoint_path, original_data)
        .await
        .unwrap();

    // Overwrite (should succeed in S3)
    backend.store(&checkpoint_path, new_data).await.unwrap();

    // Verify overwrite
    let retrieved = backend.retrieve(&checkpoint_path).await.unwrap();
    assert_eq!(retrieved, new_data);
}

#[tokio::test]
async fn test_checkpoint_metadata_size_validation() {
    let backend = create_test_backend();

    let sizes = [
        1024,              // 1KB
        1024 * 1024,       // 1MB
        10 * 1024 * 1024,  // 10MB
        100 * 1024 * 1024, // 100MB (large checkpoint)
    ];

    for (idx, &size) in sizes.iter().enumerate() {
        let path = format!("models/size_test/checkpoint_{}.bin", idx);
        let data = vec![0xFF; size];

        backend.store(&path, &data).await.unwrap();

        let metadata = backend.metadata(&path).await.unwrap();
        assert_eq!(
            metadata.size, size as u64,
            "Size mismatch for checkpoint {}",
            idx
        );
    }
}