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

//! Mock downloader implementations for error handling tests

use crate::download_types::{DownloadRequest, DownloadResult};
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::time::Duration;

// ============================================================================
// Downloader Configuration
// ============================================================================

#[derive(Debug, Clone)]
pub enum ErrorMode {
    NetworkFailure,
    RateLimited,
    InvalidAuth,
    CorruptedData,
    InvalidFormat,
    DiskFull,
    PartialFailure,
    Custom(String),
}

#[derive(Clone)]
pub struct TestDownloader {
    error_mode: Option<ErrorMode>,
    max_failures: u32,
    timeout: Option<Duration>,
    retry_delays: Arc<Mutex<Vec<Duration>>>,
    retry_count: Arc<Mutex<u32>>,
    failure_count: Arc<Mutex<u32>>,
}

impl TestDownloader {
    pub fn new() -> Self {
        Self {
            error_mode: None,
            max_failures: 0,
            timeout: None,
            retry_delays: Arc::new(Mutex::new(Vec::new())),
            retry_count: Arc::new(Mutex::new(0)),
            failure_count: Arc::new(Mutex::new(0)),
        }
    }

    pub fn with_error_mode(mut self, mode: ErrorMode) -> Self {
        self.error_mode = Some(mode);
        self
    }

    pub fn with_max_failures(mut self, max: u32) -> Self {
        self.max_failures = max;
        self
    }

    pub fn with_timeout(mut self, timeout: Duration) -> Self {
        self.timeout = Some(timeout);
        self
    }

    pub async fn download(
        &self,
        _request: DownloadRequest,
    ) -> Result<DownloadResult, Box<dyn std::error::Error>> {
        // Handle timeout mode
        if let Some(timeout) = self.timeout {
            tokio::time::sleep(timeout + Duration::from_millis(10)).await;
            return Err("Request timeout exceeded".into());
        }

        // Handle authentication errors (never retry)
        if matches!(self.error_mode, Some(ErrorMode::InvalidAuth)) {
            return Err("Authentication failed: invalid API key".into());
        }

        // Retry loop for transient errors
        let max_retries = 3;
        let mut last_error = None;

        for attempt in 0..=max_retries {
            // Check if we should fail on this attempt
            if let Some(ref mode) = self.error_mode {
                let should_fail = {
                    let count = self.failure_count.lock().expect("INVARIANT: Lock should not be poisoned");
                    *count < self.max_failures
                };

                if should_fail {
                    // Increment failure count
                    {
                        let mut count = self.failure_count.lock().expect("INVARIANT: Lock should not be poisoned");
                        *count += 1;
                    }

                    // Increment retry count
                    {
                        let mut retry_count = self.retry_count.lock().expect("INVARIANT: Lock should not be poisoned");
                        *retry_count += 1;
                    }

                    // Calculate and record exponential backoff delay
                    let base_delay = Duration::from_secs(1);
                    let delay = base_delay * 2_u32.pow(attempt);

                    // Record the delay
                    {
                        let mut delays = self.retry_delays.lock().expect("INVARIANT: Lock should not be poisoned");
                        delays.push(delay);
                    }

                    // Simulate the delay
                    tokio::time::sleep(delay).await;

                    // Store error for this attempt
                    last_error = Some(match mode {
                        ErrorMode::NetworkFailure => "Network connection failed",
                        ErrorMode::RateLimited => "Rate limit exceeded: retry after 5 seconds",
                        ErrorMode::CorruptedData => "Checksum verification failed: data corruption detected",
                        ErrorMode::InvalidFormat => "Invalid response format: failed to parse JSON",
                        ErrorMode::DiskFull => "Insufficient disk space for download",
                        ErrorMode::PartialFailure => "Download interrupted mid-transfer",
                        ErrorMode::Custom(ref msg) => msg.as_str(),
                        ErrorMode::InvalidAuth => "Authentication failed: invalid API key",
                    });

                    // Continue to next retry
                    continue;
                }
            }

            // Success case - no more failures or no error mode
            let retry_count = *self.retry_count.lock().expect("INVARIANT: Lock should not be poisoned");
            let was_rate_limited = matches!(self.error_mode, Some(ErrorMode::RateLimited));
            let total_wait_time = if was_rate_limited {
                // Calculate total wait time from retry delays
                let delays = self.retry_delays.lock().expect("INVARIANT: Lock should not be poisoned");
                delays.iter().sum::<Duration>() + Duration::from_secs(5)
            } else {
                Duration::from_secs(0)
            };

            return Ok(DownloadResult {
                retry_count,
                was_rate_limited,
                total_wait_time,
            });
        }

        // All retries exhausted
        Err(last_error.unwrap_or("All retries exhausted").into())
    }

    pub fn get_retry_delays(&self) -> Vec<Duration> {
        self.retry_delays.lock().expect("INVARIANT: Lock should not be poisoned").clone()
    }

    pub fn get_retry_count(&self) -> u32 {
        *self.retry_count.lock().expect("INVARIANT: Lock should not be poisoned")
    }
}

// ============================================================================
// Helper Functions for Error Handling Tests
// ============================================================================

pub async fn create_test_downloader_with_network_issues(_path: &Path) -> TestDownloader {
    TestDownloader::new()
        .with_error_mode(ErrorMode::NetworkFailure)
        .with_max_failures(2) // Fail twice, then succeed
}

pub async fn create_test_downloader_with_retry_tracking(_path: &Path) -> TestDownloader {
    TestDownloader::new()
        .with_error_mode(ErrorMode::NetworkFailure)
        .with_max_failures(2) // 2 retries
}

pub async fn create_test_downloader_with_rate_limiting(_path: &Path) -> TestDownloader {
    TestDownloader::new()
        .with_error_mode(ErrorMode::RateLimited)
        .with_max_failures(1) // Rate limited once, then succeed
}

pub async fn create_test_downloader_with_invalid_auth(_path: &Path) -> TestDownloader {
    TestDownloader::new()
        .with_error_mode(ErrorMode::InvalidAuth)
        .with_max_failures(1) // Auth fails immediately
}

pub async fn create_test_downloader_with_timeout(
    _path: &Path,
    timeout: Duration,
) -> TestDownloader {
    TestDownloader::new().with_timeout(timeout)
}

pub async fn create_test_downloader_with_corrupted_data(_path: &Path) -> TestDownloader {
    TestDownloader::new()
        .with_error_mode(ErrorMode::CorruptedData)
        .with_max_failures(100) // Always fail with corruption (more than max retries)
}

pub async fn create_test_downloader_with_invalid_format(_path: &Path) -> TestDownloader {
    TestDownloader::new()
        .with_error_mode(ErrorMode::InvalidFormat)
        .with_max_failures(100) // Always fail with invalid format (more than max retries)
}

pub async fn create_test_downloader_with_limited_disk(_path: &Path) -> TestDownloader {
    TestDownloader::new()
        .with_error_mode(ErrorMode::DiskFull)
        .with_max_failures(100) // Always fail with disk full (more than max retries)
}

pub async fn create_test_downloader_that_fails_midway(_path: &Path) -> TestDownloader {
    TestDownloader::new()
        .with_error_mode(ErrorMode::PartialFailure)
        .with_max_failures(100) // Always fail mid-download (more than max retries)
}

pub async fn create_test_downloader_with_error_type(
    _path: &Path,
    error_type: &str,
) -> TestDownloader {
    let error_mode = match error_type {
        "network" => ErrorMode::Custom("Failed to connect to Databento API".to_string()),
        "auth" => ErrorMode::Custom("Authentication failed: invalid API key".to_string()),
        "rate_limit" => ErrorMode::Custom("Rate limit exceeded: retry after 5 seconds".to_string()),
        "not_found" => ErrorMode::Custom("Symbol not found in dataset".to_string()),
        "invalid_date" => {
            ErrorMode::Custom("Invalid date range: start_date must be before end_date".to_string())
        },
        _ => ErrorMode::Custom(format!("Unknown error type: {}", error_type)),
    };

    TestDownloader::new()
        .with_error_mode(error_mode)
        .with_max_failures(100) // Always fail (more than max retries)
}

// ============================================================================
// Test Service with Concurrency Limit
// ============================================================================

pub struct TestService {
    concurrency_limit: usize,
    active_downloads: Arc<Mutex<Vec<String>>>,
    jobs: Arc<Mutex<std::collections::HashMap<String, crate::download_types::JobDetails>>>,
}

impl TestService {
    pub fn new(concurrency_limit: usize) -> Self {
        Self {
            concurrency_limit,
            active_downloads: Arc::new(Mutex::new(Vec::new())),
            jobs: Arc::new(Mutex::new(std::collections::HashMap::new())),
        }
    }

    pub async fn schedule_download(
        &self,
        _request: DownloadRequest,
    ) -> Result<crate::download_types::ScheduleResponse, Box<dyn std::error::Error>> {
        let job_id = uuid::Uuid::new_v4().to_string();

        // Determine initial status based on concurrency limit
        let status = {
            let active = self.active_downloads.lock().expect("INVARIANT: Lock should not be poisoned");
            if active.len() < self.concurrency_limit {
                2 // DOWNLOADING
            } else {
                1 // PENDING
            }
        };

        // Store job details
        {
            let mut jobs = self.jobs.lock().expect("INVARIANT: Lock should not be poisoned");
            jobs.insert(job_id.clone(), crate::download_types::JobDetails { status });
        }

        // Add to active downloads if downloading
        if status == 2 {
            let mut active = self.active_downloads.lock().expect("INVARIANT: Lock should not be poisoned");
            active.push(job_id.clone());
        }

        // Spawn background task to simulate download
        let jobs_clone = self.jobs.clone();
        let active_clone = self.active_downloads.clone();
        let job_id_clone = job_id.clone();
        tokio::spawn(async move {
            tokio::time::sleep(Duration::from_millis(200)).await;
            // Remove from active
            {
                let mut active = active_clone.lock().expect("INVARIANT: Lock should not be poisoned");
                active.retain(|id| id != &job_id_clone);
            }
            // Update status to completed
            {
                let mut jobs = jobs_clone.lock().expect("INVARIANT: Lock should not be poisoned");
                if let Some(job) = jobs.get_mut(&job_id_clone) {
                    job.status = 5; // COMPLETED
                }
            }
        });

        Ok(crate::download_types::ScheduleResponse { job_id })
    }

    pub async fn get_download_status(
        &self,
        job_id: String,
    ) -> Result<crate::download_types::StatusResponse, Box<dyn std::error::Error>> {
        let jobs = self.jobs.lock().expect("INVARIANT: Lock should not be poisoned");
        let job_details = jobs.get(&job_id).ok_or("Job not found")?.clone();

        Ok(crate::download_types::StatusResponse { job_details })
    }
}

pub async fn create_test_service_with_concurrency_limit(_path: &Path, limit: usize) -> TestService {
    TestService::new(limit)
}