foxhunt/services/data_acquisition_service/tests/common/mock_uploader.rs

//! Mock MinIO uploader implementation for upload tests

use crate::common::types::{ObjectMetadata, UploadResult};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::time::Duration;

// ============================================================================
// Mock Uploader State
// ============================================================================

#[derive(Clone)]
pub struct TestUploader {
    // In-memory storage for "uploaded" files
    storage: Arc<Mutex<HashMap<String, StoredObject>>>,
    // Configuration for failure simulation
    failure_count: Arc<Mutex<u32>>,
    max_failures: u32,
}

#[derive(Clone, Debug)]
struct StoredObject {
    data: Vec<u8>,
    tags: HashMap<String, String>,
    checksum: String,
}

impl TestUploader {
    pub fn new() -> Self {
        Self {
            storage: Arc::new(Mutex::new(HashMap::new())),
            failure_count: Arc::new(Mutex::new(0)),
            max_failures: 0,
        }
    }

    pub fn with_failures(max_failures: u32) -> Self {
        Self {
            storage: Arc::new(Mutex::new(HashMap::new())),
            failure_count: Arc::new(Mutex::new(0)),
            max_failures,
        }
    }

    fn should_fail(&self) -> bool {
        let mut count = self.failure_count.lock().unwrap();
        if *count < self.max_failures {
            *count += 1;
            true
        } else {
            false
        }
    }

    fn calculate_checksum(data: &[u8]) -> String {
        let mut hasher = Sha256::new();
        hasher.update(data);
        format!("{:x}", hasher.finalize())
    }

    pub async fn upload_file(
        &self,
        file_path: &Path,
        object_key: &str,
        _content_type: Option<String>,
    ) -> Result<UploadResult, Box<dyn std::error::Error + Send + Sync>> {
        // Check if file exists
        if !file_path.exists() {
            return Err("file not found".into());
        }

        // Simulate transient failures
        let mut retry_count = 0;
        while self.should_fail() {
            retry_count += 1;
            if retry_count > 3 {
                return Err("max retries exceeded".into());
            }
            tokio::time::sleep(Duration::from_millis(50)).await;
        }

        // Read file data
        let data = std::fs::read(file_path)?;
        let size_bytes = data.len() as u64;
        let checksum = Self::calculate_checksum(&data);

        // Store in mock storage
        {
            let mut storage = self.storage.lock().unwrap();
            storage.insert(
                object_key.to_string(),
                StoredObject {
                    data,
                    tags: HashMap::new(),
                    checksum: checksum.clone(),
                },
            );
        }

        Ok(UploadResult {
            object_url: format!("s3://test-bucket/{}", object_key),
            size_bytes,
            upload_duration_ms: 100,
            retry_count,
            checksum,
        })
    }

    pub async fn upload_file_with_tags(
        &self,
        file_path: &Path,
        object_key: &str,
        content_type: Option<String>,
        tags: HashMap<String, String>,
    ) -> Result<UploadResult, Box<dyn std::error::Error + Send + Sync>> {
        // Upload file first
        let result = self.upload_file(file_path, object_key, content_type).await?;

        // Store tags
        {
            let mut storage = self.storage.lock().unwrap();
            if let Some(obj) = storage.get_mut(object_key) {
                obj.tags = tags;
            }
        }

        Ok(result)
    }

    pub async fn upload_file_with_progress<F>(
        &self,
        file_path: &Path,
        object_key: &str,
        content_type: Option<String>,
        callback: F,
    ) -> Result<UploadResult, Box<dyn std::error::Error + Send + Sync>>
    where
        F: Fn(u64, u64) + Send + 'static,
    {
        // Check if file exists
        if !file_path.exists() {
            return Err("file not found".into());
        }

        // Get file size
        let file_size = std::fs::metadata(file_path)?.len();
        let chunk_size = 1024 * 1024; // 1 MB chunks

        // Simulate chunked upload with progress callbacks
        let mut uploaded = 0u64;
        while uploaded < file_size {
            tokio::time::sleep(Duration::from_millis(10)).await;

            uploaded = std::cmp::min(uploaded + chunk_size, file_size);

            // Invoke progress callback
            callback(uploaded, file_size);
        }

        // Perform actual upload
        self.upload_file(file_path, object_key, content_type).await
    }

    pub async fn get_object_metadata(
        &self,
        object_key: &str,
    ) -> Result<ObjectMetadata, Box<dyn std::error::Error + Send + Sync>> {
        let storage = self.storage.lock().unwrap();
        let obj = storage
            .get(object_key)
            .ok_or("Object not found")?;

        Ok(ObjectMetadata {
            tags: obj.tags.clone(),
        })
    }
}

// ============================================================================
// Helper Functions for MinIO Upload Tests
// ============================================================================

pub async fn create_test_uploader() -> TestUploader {
    TestUploader::new()
}

pub async fn create_test_uploader_with_failures(num_failures: u32) -> TestUploader {
    TestUploader::with_failures(num_failures)
}