The storage crate had its own RetryConfig struct (max_attempts, initial_delay, max_delay, backoff_multiplier) duplicating common::resilience::retry::RetryConfig. Added backoff_multiplier field to common's RetryConfig and updated storage to re-export and use common's version with its field names (max_retries, base_delay). Updated all storage tests accordingly. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
483 lines
13 KiB
Rust
483 lines
13 KiB
Rust
//! 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 object_store::ObjectStore;
|
|
use storage::model_helpers::{ConnectionPool, ProgressCallback, RetryConfig};
|
|
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,
|
|
"test-bucket".to_string(),
|
|
)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_store_and_retrieve() {
|
|
let backend = create_test_backend();
|
|
|
|
// 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();
|
|
|
|
// 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();
|
|
|
|
// 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();
|
|
|
|
// Delete nonexistent file
|
|
// Note: InMemory store returns Ok(()) even for non-existent files,
|
|
// so this tests the happy path rather than the NotFound case
|
|
let result = backend.delete("nonexistent.txt").await;
|
|
assert!(
|
|
result.is_ok(),
|
|
"delete should not fail for non-existent file"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_list() {
|
|
let backend = create_test_backend();
|
|
|
|
// 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();
|
|
|
|
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();
|
|
|
|
// Create connection pool with properly typed stores
|
|
let store1: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
|
|
let store2: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
|
|
|
|
let pool = Arc::new(ConnectionPool::new(vec![store1, store2]));
|
|
|
|
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();
|
|
|
|
let retry_config = RetryConfig {
|
|
max_retries: 5,
|
|
base_delay: std::time::Duration::from_millis(50),
|
|
max_delay: std::time::Duration::from_secs(1),
|
|
backoff_multiplier: 1.5,
|
|
..Default::default()
|
|
};
|
|
|
|
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();
|
|
|
|
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();
|
|
|
|
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();
|
|
|
|
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();
|
|
|
|
// 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();
|
|
|
|
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();
|
|
|
|
// 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();
|
|
|
|
// 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();
|
|
|
|
// 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();
|
|
|
|
// 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();
|
|
|
|
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();
|
|
|
|
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();
|
|
|
|
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();
|
|
|
|
// 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();
|
|
|
|
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();
|
|
|
|
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());
|
|
|
|
let mut handles = vec![];
|
|
|
|
for i in 0..10 {
|
|
let backend = Arc::clone(&backend);
|
|
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 - list all and filter
|
|
let files = backend.list("").await.unwrap();
|
|
let concurrent_files: Vec<_> = files.iter().filter(|f| f.contains("concurrent_")).collect();
|
|
assert_eq!(concurrent_files.len(), 10);
|
|
}
|