//! 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, 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, } 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>>, 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> { 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>>, 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> { 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, _start_time: Option, _end_time: Option, ) -> Result> { 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> { 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) }