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

//! Mock data acquisition service implementation for workflow tests

use crate::workflow_types::{
    CancelDownloadResponse, DownloadJobDetails, GetDownloadStatusResponse,
    ListDownloadJobsResponse, ScheduleDownloadRequest, ScheduleDownloadResponse,
};
use std::collections::HashMap;
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::time::Duration;

// ============================================================================
// Download Status Constants (matching proto enum)
// ============================================================================

const STATUS_PENDING: i32 = 1;
const STATUS_DOWNLOADING: i32 = 2;
const STATUS_VALIDATING: i32 = 3;
const STATUS_UPLOADING: i32 = 4;
const STATUS_COMPLETED: i32 = 5;
const STATUS_FAILED: i32 = 6;
const STATUS_CANCELLED: i32 = 7;

// ============================================================================
// Mock Job State
// ============================================================================

#[derive(Clone, Debug)]
struct JobState {
    job_id: String,
    status: i32,
    dataset: String,
    symbols: Vec<String>,
    start_date: String,
    end_date: String,
    progress_percentage: f32,
    created_at: i64,
    completed_at: i64,
    minio_path: String,
    records_count: u64,
    data_quality_score: f64,
    invalid_records: u64,
    cancellation_reason: Option<String>,
}

impl JobState {
    fn new(job_id: String, request: ScheduleDownloadRequest) -> Self {
        Self {
            job_id,
            status: STATUS_PENDING,
            dataset: request.dataset,
            symbols: request.symbols,
            start_date: request.start_date,
            end_date: request.end_date,
            progress_percentage: 0.0,
            created_at: chrono::Utc::now().timestamp(),
            completed_at: 0,
            minio_path: String::new(),
            records_count: 0,
            data_quality_score: 1.0,
            invalid_records: 0,
            cancellation_reason: None,
        }
    }

    fn estimate_cost(start_date: &str, end_date: &str, symbols: &[String]) -> f64 {
        // Parse dates and calculate days
        let start = chrono::NaiveDate::parse_from_str(start_date, "%Y-%m-%d")
            .unwrap_or_else(|_| chrono::NaiveDate::from_ymd_opt(2024, 1, 1).unwrap());
        let end = chrono::NaiveDate::parse_from_str(end_date, "%Y-%m-%d")
            .unwrap_or_else(|_| chrono::NaiveDate::from_ymd_opt(2024, 1, 2).unwrap());

        let days = (end - start).num_days().max(1) as f64;
        let num_symbols = symbols.len() as f64;

        // Simple cost model: $1 per symbol per day
        days * num_symbols
    }

    fn to_job_details(&self) -> DownloadJobDetails {
        DownloadJobDetails {
            job_id: self.job_id.clone(),
            status: self.status,
            dataset: self.dataset.clone(),
            symbols: self.symbols.clone(),
            progress_percentage: self.progress_percentage,
            completed_at: self.completed_at,
            minio_path: self.minio_path.clone(),
            records_count: self.records_count,
            data_quality_score: self.data_quality_score,
            invalid_records: self.invalid_records,
        }
    }
}

// ============================================================================
// Mock Service Implementation
// ============================================================================

#[derive(Clone)]
pub struct TestDataAcquisitionService {
    jobs: Arc<Mutex<HashMap<String, JobState>>>,
    simulate_corrupted_data: bool,
}

impl TestDataAcquisitionService {
    pub fn new(simulate_corrupted_data: bool) -> Self {
        Self {
            jobs: Arc::new(Mutex::new(HashMap::new())),
            simulate_corrupted_data,
        }
    }

    pub async fn schedule_download(
        &self,
        request: ScheduleDownloadRequest,
    ) -> Result<ScheduleDownloadResponse, Box<dyn std::error::Error>> {
        let job_id = uuid::Uuid::new_v4().to_string();
        let estimated_cost =
            JobState::estimate_cost(&request.start_date, &request.end_date, &request.symbols);

        let job_state = JobState::new(job_id.clone(), request);

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

        // Spawn background task to progress job through states
        let jobs_clone = self.jobs.clone();
        let job_id_clone = job_id.clone();
        let simulate_corrupted = self.simulate_corrupted_data;

        tokio::spawn(async move {
            Self::progress_job_states(jobs_clone, job_id_clone, simulate_corrupted).await;
        });

        Ok(ScheduleDownloadResponse {
            job_id,
            status: STATUS_PENDING,
            estimated_cost_usd: estimated_cost,
        })
    }

    async fn progress_job_states(
        jobs: Arc<Mutex<HashMap<String, JobState>>>,
        job_id: String,
        simulate_corrupted: bool,
    ) {
        let states = vec![
            (STATUS_DOWNLOADING, 25.0, 100),
            (STATUS_VALIDATING, 50.0, 150),
            (STATUS_UPLOADING, 75.0, 100),
            (STATUS_COMPLETED, 100.0, 50),
        ];

        for (status, progress, delay_ms) in states {
            tokio::time::sleep(Duration::from_millis(delay_ms)).await;

            let mut jobs = jobs.lock().expect("INVARIANT: Lock should not be poisoned");
            if let Some(job) = jobs.get_mut(&job_id) {
                // Check if job was cancelled
                if job.status == STATUS_CANCELLED {
                    return;
                }

                job.status = status;
                job.progress_percentage = progress;

                // Simulate data quality issues for corrupted data test
                if simulate_corrupted && status == STATUS_VALIDATING {
                    job.data_quality_score = 0.85; // Low quality
                    job.invalid_records = 150; // Some invalid records
                }

                // Set completion metadata
                if status == STATUS_COMPLETED {
                    job.completed_at = chrono::Utc::now().timestamp();
                    job.minio_path = format!(
                        "market-data/{}/{}_{}.dbn",
                        job.symbols.join("-"),
                        job.start_date,
                        job.end_date
                    );
                    job.records_count = 10000;

                    // Use good quality if not corrupted
                    if !simulate_corrupted {
                        job.data_quality_score = 0.99;
                        job.invalid_records = 10;
                    }
                }
            }
        }
    }

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

        Ok(GetDownloadStatusResponse {
            job_details: job.to_job_details(),
        })
    }

    pub async fn list_download_jobs(
        &self,
        page: u32,
        page_size: u32,
        status_filter: Option<i32>,
        _start_time: Option<i64>,
        _end_time: Option<i64>,
    ) -> Result<ListDownloadJobsResponse, Box<dyn std::error::Error>> {
        let jobs = self.jobs.lock().expect("INVARIANT: Lock should not be poisoned");

        // Collect and filter jobs
        let mut all_jobs: Vec<_> = jobs.values().cloned().collect();

        // Apply status filter if provided
        if let Some(status) = status_filter {
            all_jobs.retain(|job| job.status == status);
        }

        // Sort by created_at (newest first)
        all_jobs.sort_by(|a, b| b.created_at.cmp(&a.created_at));

        let total_count = all_jobs.len() as u32;

        // Apply pagination
        let start = ((page - 1) * page_size) as usize;
        let end = (start + page_size as usize).min(all_jobs.len());
        let paginated_jobs: Vec<_> = all_jobs[start..end]
            .iter()
            .map(|job| job.to_job_details())
            .collect();

        Ok(ListDownloadJobsResponse {
            jobs: paginated_jobs,
            total_count,
            page,
            page_size,
        })
    }

    pub async fn cancel_download(
        &self,
        job_id: String,
        reason: String,
    ) -> Result<CancelDownloadResponse, Box<dyn std::error::Error>> {
        let mut jobs = self.jobs.lock().expect("INVARIANT: Lock should not be poisoned");
        let job = jobs.get_mut(&job_id).ok_or("Job not found")?;

        // Only cancel if not already completed or failed
        if job.status != STATUS_COMPLETED && job.status != STATUS_FAILED {
            job.status = STATUS_CANCELLED;
            job.cancellation_reason = Some(reason);
            Ok(CancelDownloadResponse { success: true })
        } else {
            Err("Cannot cancel completed or failed job".into())
        }
    }
}

// ============================================================================
// Helper Functions for Download Workflow Tests
// ============================================================================

pub async fn create_test_service(_path: &Path) -> TestDataAcquisitionService {
    TestDataAcquisitionService::new(false)
}

pub async fn create_test_service_with_corrupted_data(_path: &Path) -> TestDataAcquisitionService {
    TestDataAcquisitionService::new(true)
}