foxhunt/storage/tests/object_store_backend_tests.rs

//! Comprehensive tests for ObjectStoreBackend
//!
//! Tests S3 backend using in-memory object store to avoid network dependencies.

use std::sync::Arc;

use object_store::memory::InMemory;
use storage::object_store_backend::ObjectStoreBackend;
use storage::model_helpers::{ConnectionPool, ProgressCallback, RetryConfig};
use storage::Storage;

/// Helper to create test backend with in-memory store
async fn create_test_backend() -> ObjectStoreBackend {
    let config = config::schemas::S3Config {
        bucket_name: "test-bucket".to_string(),
        region: "us-east-1".to_string(),
        access_key_id: Some("test_key".to_string()),
        secret_access_key: Some("test_secret".to_string()),
        session_token: None,
        endpoint_url: None,
        force_path_style: false,
        timeout: std::time::Duration::from_secs(30),
        max_retry_attempts: 3,
        use_ssl: true,
    };

    // Create backend with in-memory store for testing
    let in_memory_store = Arc::new(InMemory::new());
    ObjectStoreBackend {
        store: in_memory_store,
        pool: None,
        _bucket: "test-bucket".to_string(),
        _config: config,
        retry_config: RetryConfig::default(),
    }
}

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

    // Store data
    let test_data = b"test data for storage";
    backend
        .store("test/file.bin", test_data)
        .await
        .expect("Failed to store");

    // Retrieve and verify
    let retrieved = backend
        .retrieve("test/file.bin")
        .await
        .expect("Failed to retrieve");
    assert_eq!(retrieved, test_data);
}

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

    // Should not exist initially
    let exists = backend.exists("nonexistent.txt").await.expect("exists failed");
    assert!(!exists);

    // Store file
    backend
        .store("exists.txt", b"data")
        .await
        .expect("store failed");

    // Should exist now
    let exists = backend.exists("exists.txt").await.expect("exists failed");
    assert!(exists);
}

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

    // Store file
    backend
        .store("delete_me.txt", b"data")
        .await
        .expect("store failed");

    // Verify exists
    assert!(backend.exists("delete_me.txt").await.expect("exists failed"));

    // Delete
    let deleted = backend
        .delete("delete_me.txt")
        .await
        .expect("delete failed");
    assert!(deleted);

    // Verify gone
    assert!(!backend.exists("delete_me.txt").await.expect("exists failed"));
}

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

    // Delete nonexistent file should return false
    let deleted = backend
        .delete("nonexistent.txt")
        .await
        .expect("delete failed");
    assert!(!deleted);
}

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

    // Store multiple files
    backend.store("models/v1/weights.bin", b"data1").await.unwrap();
    backend.store("models/v1/config.json", b"data2").await.unwrap();
    backend.store("models/v2/weights.bin", b"data3").await.unwrap();
    backend.store("data/test.csv", b"data4").await.unwrap();

    // List with prefix
    let files = backend.list("models/v1/").await.expect("list failed");
    assert_eq!(files.len(), 2);
    assert!(files.iter().any(|f| f.contains("weights.bin")));
    assert!(files.iter().any(|f| f.contains("config.json")));

    // List all models
    let files = backend.list("models/").await.expect("list failed");
    assert_eq!(files.len(), 3);

    // List all
    let files = backend.list("").await.expect("list failed");
    assert_eq!(files.len(), 4);
}

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

    let test_data = b"test metadata";
    backend
        .store("metadata_test.txt", test_data)
        .await
        .unwrap();

    let metadata = backend
        .metadata("metadata_test.txt")
        .await
        .expect("metadata failed");

    assert_eq!(metadata.path, "metadata_test.txt");
    assert_eq!(metadata.size, test_data.len() as u64);
    assert!(metadata.etag.is_some());
}

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

    // Create connection pool with multiple stores
    let store1 = Arc::new(InMemory::new());
    let store2 = Arc::new(InMemory::new());
    let pool = Arc::new(ConnectionPool::new(vec![store1.clone(), store2.clone()]));

    let backend = backend.with_connection_pool(pool);

    // Store data
    backend.store("pool_test.txt", b"data").await.unwrap();

    // Retrieve
    let data = backend.retrieve("pool_test.txt").await.unwrap();
    assert_eq!(data, b"data");
}

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

    let retry_config = RetryConfig {
        max_attempts: 5,
        initial_delay: std::time::Duration::from_millis(50),
        max_delay: std::time::Duration::from_secs(1),
        backoff_multiplier: 1.5,
    };

    let backend = backend.with_retry_config(retry_config);

    // Test operation succeeds with new retry config
    backend.store("retry_test.txt", b"data").await.unwrap();
    let data = backend.retrieve("retry_test.txt").await.unwrap();
    assert_eq!(data, b"data");
}

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

    let path = backend.get_model_path("mamba", "v1.0", "weights.safetensors");
    assert_eq!(path, "models/mamba/v1.0/weights.safetensors");
}

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

    let path = backend.get_checkpoint_path("dqn", "epoch_100");
    assert_eq!(path, "models/dqn/checkpoints/epoch_100");
}

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

    let path = backend.get_metadata_path("ppo", "v2.1");
    assert_eq!(path, "models/ppo/v2.1/metadata.json");
}

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

    // Store test file
    let test_data = vec![0u8; 1024]; // 1KB
    backend
        .store("progress_test.bin", &test_data)
        .await
        .unwrap();

    // Track progress callbacks
    let progress_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
    let progress_count_clone = progress_count.clone();

    let callback: ProgressCallback = Arc::new(move |downloaded, total| {
        progress_count_clone.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        println!("Progress: {}/{} bytes", downloaded, total);
    });

    // Download with progress
    let data = backend
        .download_with_progress("progress_test.bin", Some(callback))
        .await
        .expect("download failed");

    assert_eq!(data.len(), test_data.len());
    // Should have at least 2 callbacks (initial + final)
    assert!(
        progress_count.load(std::sync::atomic::Ordering::SeqCst) >= 2
    );
}

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

    let test_data = b"test data";
    backend.store("no_callback.txt", test_data).await.unwrap();

    // Download without callback
    let data = backend
        .download_with_progress("no_callback.txt", None)
        .await
        .expect("download failed");

    assert_eq!(data, test_data);
}

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

    // Store larger test file
    let test_data = vec![42u8; 4096]; // 4KB
    backend
        .store("stream_test.bin", &test_data)
        .await
        .unwrap();

    // Track progress
    let progress_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
    let progress_count_clone = progress_count.clone();

    let callback: ProgressCallback = Arc::new(move |downloaded, total| {
        progress_count_clone.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        assert!(downloaded <= total);
    });

    // Stream download
    let data = backend
        .stream_download_with_progress("stream_test.bin", 512, callback)
        .await
        .expect("stream download failed");

    assert_eq!(data.len(), test_data.len());
    // Should have multiple progress callbacks for streaming
    assert!(
        progress_count.load(std::sync::atomic::Ordering::SeqCst) >= 1
    );
}

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

    // Store multiple files
    backend.store("file1.txt", b"data1").await.unwrap();
    backend.store("file2.txt", b"data2").await.unwrap();
    backend.store("file3.txt", b"data3").await.unwrap();

    let paths = vec![
        "file1.txt".to_string(),
        "file2.txt".to_string(),
        "file3.txt".to_string(),
    ];

    // Parallel download without pool (falls back to sequential)
    let results = backend
        .parallel_download(paths, None)
        .await
        .expect("parallel download failed");

    assert_eq!(results.len(), 3);
    assert_eq!(results[0].1, b"data1");
    assert_eq!(results[1].1, b"data2");
    assert_eq!(results[2].1, b"data3");
}

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

    // Store files
    backend.store("p1.txt", b"data1").await.unwrap();
    backend.store("p2.txt", b"data2").await.unwrap();

    let paths = vec!["p1.txt".to_string(), "p2.txt".to_string()];

    let progress_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
    let progress_count_clone = progress_count.clone();

    let callback: ProgressCallback = Arc::new(move |current, total| {
        progress_count_clone.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        assert!(current <= total);
    });

    let results = backend
        .parallel_download(paths, Some(callback))
        .await
        .expect("parallel download failed");

    assert_eq!(results.len(), 2);
    assert!(
        progress_count.load(std::sync::atomic::Ordering::SeqCst) >= 2
    );
}

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

    // Store 1MB file
    let large_data = vec![0xAAu8; 1_048_576];
    backend
        .store("large_file.bin", &large_data)
        .await
        .expect("store large file failed");

    // Verify size via metadata
    let metadata = backend.metadata("large_file.bin").await.unwrap();
    assert_eq!(metadata.size, 1_048_576);

    // Retrieve and verify
    let retrieved = backend.retrieve("large_file.bin").await.unwrap();
    assert_eq!(retrieved.len(), large_data.len());
}

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

    let nested_path = "models/mamba/v1.0/checkpoints/epoch_100/weights.safetensors";
    backend.store(nested_path, b"weights").await.unwrap();

    let data = backend.retrieve(nested_path).await.unwrap();
    assert_eq!(data, b"weights");

    assert!(backend.exists(nested_path).await.unwrap());
}

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

    let paths = vec![
        "test-file.txt",
        "test_file_2.txt",
        "test.file.3.txt",
        "models/v1.0/weights.safetensors",
    ];

    for path in paths {
        backend.store(path, b"data").await.unwrap();
        let data = backend.retrieve(path).await.unwrap();
        assert_eq!(data, b"data");
    }
}

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

    backend.store("empty.txt", b"").await.unwrap();

    let data = backend.retrieve("empty.txt").await.unwrap();
    assert_eq!(data, b"");

    let metadata = backend.metadata("empty.txt").await.unwrap();
    assert_eq!(metadata.size, 0);
}

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

    // Store initial data
    backend.store("overwrite.txt", b"initial").await.unwrap();
    let data = backend.retrieve("overwrite.txt").await.unwrap();
    assert_eq!(data, b"initial");

    // Overwrite
    backend.store("overwrite.txt", b"updated").await.unwrap();
    let data = backend.retrieve("overwrite.txt").await.unwrap();
    assert_eq!(data, b"updated");
}

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

    backend.store("file1.txt", b"data1").await.unwrap();
    backend.store("file2.txt", b"data2").await.unwrap();

    // List with empty prefix should return all files
    let files = backend.list("").await.unwrap();
    assert_eq!(files.len(), 2);
}

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

    backend.store("models/v1/weights.bin", b"data").await.unwrap();

    // List with non-matching prefix
    let files = backend.list("models/v2/").await.unwrap();
    assert_eq!(files.len(), 0);
}

#[tokio::test]
async fn test_concurrent_operations() {
    let backend = Arc::new(create_test_backend().await);

    let mut handles = vec![];

    for i in 0..10 {
        let backend = backend.clone();
        let handle = tokio::spawn(async move {
            let path = format!("concurrent_{}.txt", i);
            let data = format!("data_{}", i);
            backend.store(&path, data.as_bytes()).await.unwrap();
            let retrieved = backend.retrieve(&path).await.unwrap();
            assert_eq!(retrieved, data.as_bytes());
        });
        handles.push(handle);
    }

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

    // Verify all files exist
    let files = backend.list("concurrent_").await.unwrap();
    assert_eq!(files.len(), 10);
}