foxhunt/bin/fxt/tests/e2e/mock_ml_training_service.rs

//! Mock ML Training Service for E2E Testing
//!
//! Provides an in-memory gRPC server that simulates the ML Training Service
//! behavior for testing TLI commands end-to-end.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use tonic::{Request, Response, Status};
use fxt::proto::ml_training::{
    ml_training_service_server::MlTrainingService, HealthCheckRequest, HealthCheckResponse,
    ListAvailableModelsRequest, ListAvailableModelsResponse, ListTrainingJobsRequest,
    ListTrainingJobsResponse, GetTrainingJobDetailsRequest, GetTrainingJobDetailsResponse,
    StartTrainingRequest, StartTrainingResponse, SubscribeToTrainingStatusRequest,
    TrainingStatusUpdate, StopTrainingRequest, StopTrainingResponse, TrainingJobSummary,
    TrainingStatus, StartTuningJobRequest, StartTuningJobResponse, GetTuningJobStatusRequest,
    GetTuningJobStatusResponse, StopTuningJobRequest, StopTuningJobResponse,
    TrainModelRequest, TrainModelResponse, StreamProgressRequest, ProgressUpdate,
    BatchStartTuningJobsRequest, BatchStartTuningJobsResponse,
    GetBatchTuningStatusRequest, GetBatchTuningStatusResponse,
    StopBatchTuningJobRequest, StopBatchTuningJobResponse,
    JobCompletionReport, JobCompletionAck,
    ListPendingPromotionsRequest, ListPendingPromotionsResponse,
    ApprovePromotionRequest, ApprovePromotionResponse,
    RejectPromotionRequest, RejectPromotionResponse,
    ApproveModelRequest, ApproveModelResponse,
    RejectModelRequest, RejectModelResponse,
};

/// Shared state for mock training service
#[derive(Debug, Clone)]
pub struct MockTrainingState {
    jobs: Arc<Mutex<HashMap<String, TrainingJobSummary>>>,
    next_job_id: Arc<Mutex<u64>>,
}

impl Default for MockTrainingState {
    fn default() -> Self {
        Self {
            jobs: Arc::new(Mutex::new(HashMap::new())),
            next_job_id: Arc::new(Mutex::new(1)),
        }
    }
}

impl MockTrainingState {
    /// Add a job to the mock state
    pub fn add_job(&self, job: TrainingJobSummary) {
        let mut jobs = self.jobs.lock().unwrap();
        jobs.insert(job.job_id.clone(), job);
    }

    /// Get a job by ID
    pub fn get_job(&self, job_id: &str) -> Option<TrainingJobSummary> {
        let jobs = self.jobs.lock().unwrap();
        jobs.get(job_id).cloned()
    }

    /// Update job status
    pub fn update_job_status(&self, job_id: &str, status: TrainingStatus) {
        let mut jobs = self.jobs.lock().unwrap();
        if let Some(job) = jobs.get_mut(job_id) {
            job.status = status as i32;
        }
    }

    /// List all jobs
    pub fn list_jobs(&self) -> Vec<TrainingJobSummary> {
        let jobs = self.jobs.lock().unwrap();
        jobs.values().cloned().collect()
    }

    /// Generate next job ID
    pub fn next_job_id(&self) -> String {
        let mut next_id = self.next_job_id.lock().unwrap();
        let id = *next_id;
        *next_id += 1;
        format!("train_job_{}", id)
    }

    /// Clear all jobs
    pub fn clear(&self) {
        let mut jobs = self.jobs.lock().unwrap();
        jobs.clear();
    }
}

/// Mock ML Training Service implementation
pub struct MockMlTrainingService {
    state: MockTrainingState,
}

impl MockMlTrainingService {
    pub fn new(state: MockTrainingState) -> Self {
        Self { state }
    }
}

#[tonic::async_trait]
impl MlTrainingService for MockMlTrainingService {
    async fn start_training(
        &self,
        request: Request<StartTrainingRequest>,
    ) -> Result<Response<StartTrainingResponse>, Status> {
        let req = request.into_inner();
        let job_id = self.state.next_job_id();

        // Create a new job
        let job = TrainingJobSummary {
            job_id: job_id.clone(),
            model_type: req.model_type.clone(),
            status: TrainingStatus::Pending as i32,
            created_at: chrono::Utc::now().timestamp(),
            started_at: 0,
            completed_at: 0,
            description: req.description.clone(),
            final_loss: 0.0,
            best_validation_score: 0.0,
            tags: req.tags.clone(),
        };

        self.state.add_job(job);

        Ok(Response::new(StartTrainingResponse {
            job_id,
            status: TrainingStatus::Pending as i32,
            message: "Training job created successfully".to_string(),
        }))
    }

    type SubscribeToTrainingStatusStream = ReceiverStream<Result<TrainingStatusUpdate, Status>>;

    async fn subscribe_to_training_status(
        &self,
        request: Request<SubscribeToTrainingStatusRequest>,
    ) -> Result<Response<Self::SubscribeToTrainingStatusStream>, Status> {
        let req = request.into_inner();
        let job_id = req.job_id.clone();

        // Check if job exists
        if self.state.get_job(&job_id).is_none() {
            return Err(Status::not_found(format!("Job {} not found", job_id)));
        }

        let (tx, rx) = mpsc::channel(100);
        let state = self.state.clone();

        // Spawn a task to simulate training progress
        tokio::spawn(async move {
            // Simulate training progress from 0% to 100%
            for epoch in 0..=10 {
                let progress = (epoch as f32 / 10.0) * 100.0;
                let status = if epoch == 0 {
                    TrainingStatus::Running
                } else if epoch == 10 {
                    TrainingStatus::Completed
                } else {
                    TrainingStatus::Running
                };

                // Update state
                state.update_job_status(&job_id, status);

                let update = TrainingStatusUpdate {
                    job_id: job_id.clone(),
                    status: status as i32,
                    progress_percentage: progress,
                    current_epoch: epoch,
                    total_epochs: 10,
                    metrics: HashMap::new(),
                    message: format!("Training epoch {}/10", epoch),
                    timestamp: chrono::Utc::now().timestamp(),
                    financial_metrics: None,
                    resource_usage: None,
                };

                if tx.send(Ok(update)).await.is_err() {
                    break;
                }

                tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
            }
        });

        Ok(Response::new(ReceiverStream::new(rx)))
    }

    async fn stop_training(
        &self,
        request: Request<StopTrainingRequest>,
    ) -> Result<Response<StopTrainingResponse>, Status> {
        let req = request.into_inner();

        // Check if job exists
        if self.state.get_job(&req.job_id).is_none() {
            return Err(Status::not_found(format!("Job {} not found", req.job_id)));
        }

        // Update job status to stopped
        self.state.update_job_status(&req.job_id, TrainingStatus::Stopped);

        Ok(Response::new(StopTrainingResponse {
            success: true,
            message: format!("Training job {} stopped successfully", req.job_id),
        }))
    }

    async fn list_available_models(
        &self,
        _request: Request<ListAvailableModelsRequest>,
    ) -> Result<Response<ListAvailableModelsResponse>, Status> {
        Ok(Response::new(ListAvailableModelsResponse {
            models: vec![],
        }))
    }

    async fn list_training_jobs(
        &self,
        request: Request<ListTrainingJobsRequest>,
    ) -> Result<Response<ListTrainingJobsResponse>, Status> {
        let req = request.into_inner();
        let mut jobs = self.state.list_jobs();

        // Apply status filter
        if req.status_filter != 0 {
            jobs.retain(|j| j.status == req.status_filter);
        }

        // Apply model type filter
        if !req.model_type_filter.is_empty() {
            jobs.retain(|j| j.model_type == req.model_type_filter);
        }

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

        let total_count = jobs.len();
        let page_size = if req.page_size == 0 { 50 } else { req.page_size } as usize;
        let page = if req.page == 0 { 1 } else { req.page } as usize;
        let start = (page - 1) * page_size;
        let end = std::cmp::min(start + page_size, jobs.len());

        let page_jobs = if start < jobs.len() {
            jobs[start..end].to_vec()
        } else {
            vec![]
        };

        Ok(Response::new(ListTrainingJobsResponse {
            jobs: page_jobs,
            total_count: total_count as u32,
            page: page as u32,
            page_size: page_size as u32,
        }))
    }

    async fn get_training_job_details(
        &self,
        request: Request<GetTrainingJobDetailsRequest>,
    ) -> Result<Response<GetTrainingJobDetailsResponse>, Status> {
        let req = request.into_inner();

        if self.state.get_job(&req.job_id).is_none() {
            return Err(Status::not_found(format!("Job {} not found", req.job_id)));
        }

        Ok(Response::new(GetTrainingJobDetailsResponse {
            job_details: None,
        }))
    }

    async fn health_check(
        &self,
        _request: Request<HealthCheckRequest>,
    ) -> Result<Response<HealthCheckResponse>, Status> {
        Ok(Response::new(HealthCheckResponse {
            healthy: true,
            message: "Mock ML Training Service is healthy".to_string(),
            details: HashMap::new(),
        }))
    }

    async fn start_tuning_job(
        &self,
        _request: Request<StartTuningJobRequest>,
    ) -> Result<Response<StartTuningJobResponse>, Status> {
        unimplemented!("Tuning not implemented in mock")
    }

    async fn get_tuning_job_status(
        &self,
        _request: Request<GetTuningJobStatusRequest>,
    ) -> Result<Response<GetTuningJobStatusResponse>, Status> {
        unimplemented!("Tuning not implemented in mock")
    }

    async fn stop_tuning_job(
        &self,
        _request: Request<StopTuningJobRequest>,
    ) -> Result<Response<StopTuningJobResponse>, Status> {
        unimplemented!("Tuning not implemented in mock")
    }

    async fn train_model(
        &self,
        _request: Request<TrainModelRequest>,
    ) -> Result<Response<TrainModelResponse>, Status> {
        unimplemented!("Train model not implemented in mock")
    }

    type StreamTuningProgressStream = ReceiverStream<Result<ProgressUpdate, Status>>;

    async fn stream_tuning_progress(
        &self,
        _request: Request<StreamProgressRequest>,
    ) -> Result<Response<Self::StreamTuningProgressStream>, Status> {
        unimplemented!("Tuning progress not implemented in mock")
    }

    async fn batch_start_tuning_jobs(
        &self,
        _request: Request<BatchStartTuningJobsRequest>,
    ) -> Result<Response<BatchStartTuningJobsResponse>, Status> {
        unimplemented!("Batch tuning not implemented in mock")
    }

    async fn get_batch_tuning_status(
        &self,
        _request: Request<GetBatchTuningStatusRequest>,
    ) -> Result<Response<GetBatchTuningStatusResponse>, Status> {
        unimplemented!("Batch tuning not implemented in mock")
    }

    async fn stop_batch_tuning_job(
        &self,
        _request: Request<StopBatchTuningJobRequest>,
    ) -> Result<Response<StopBatchTuningJobResponse>, Status> {
        unimplemented!("Batch tuning not implemented in mock")
    }

    async fn report_job_completion(
        &self,
        _request: Request<JobCompletionReport>,
    ) -> Result<Response<JobCompletionAck>, Status> {
        Ok(Response::new(JobCompletionAck {
            accepted: true,
            promotion_status: "registered".to_string(),
        }))
    }

    async fn list_pending_promotions(
        &self,
        _request: Request<ListPendingPromotionsRequest>,
    ) -> Result<Response<ListPendingPromotionsResponse>, Status> {
        Ok(Response::new(ListPendingPromotionsResponse {
            promotions: vec![],
        }))
    }

    async fn approve_promotion(
        &self,
        _request: Request<ApprovePromotionRequest>,
    ) -> Result<Response<ApprovePromotionResponse>, Status> {
        Ok(Response::new(ApprovePromotionResponse {
            success: true,
            message: "Mock approved".to_string(),
        }))
    }

    async fn reject_promotion(
        &self,
        _request: Request<RejectPromotionRequest>,
    ) -> Result<Response<RejectPromotionResponse>, Status> {
        Ok(Response::new(RejectPromotionResponse {
            success: true,
            message: "Mock rejected".to_string(),
        }))
    }

    async fn approve_model(
        &self,
        _request: Request<ApproveModelRequest>,
    ) -> Result<Response<ApproveModelResponse>, Status> {
        Ok(Response::new(ApproveModelResponse {
            success: true,
            message: "Mock model approved".to_string(),
        }))
    }

    async fn reject_model(
        &self,
        _request: Request<RejectModelRequest>,
    ) -> Result<Response<RejectModelResponse>, Status> {
        Ok(Response::new(RejectModelResponse {
            success: true,
            message: "Mock model rejected".to_string(),
        }))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_mock_state_creation() {
        let state = MockTrainingState::default();
        assert_eq!(state.list_jobs().len(), 0);
    }

    #[test]
    fn test_add_and_get_job() {
        let state = MockTrainingState::default();
        let job = TrainingJobSummary {
            job_id: "test_job_1".to_string(),
            model_type: "TFT".to_string(),
            status: TrainingStatus::Pending as i32,
            created_at: 1729602000,
            started_at: 0,
            completed_at: 0,
            description: "Test job".to_string(),
            final_loss: 0.0,
            best_validation_score: 0.0,
            tags: HashMap::new(),
        };

        state.add_job(job.clone());
        let retrieved = state.get_job("test_job_1");
        assert!(retrieved.is_some());
        assert_eq!(retrieved.unwrap().job_id, "test_job_1");
    }

    #[test]
    fn test_update_job_status() {
        let state = MockTrainingState::default();
        let job = TrainingJobSummary {
            job_id: "test_job_2".to_string(),
            model_type: "DQN".to_string(),
            status: TrainingStatus::Pending as i32,
            created_at: 1729602000,
            started_at: 0,
            completed_at: 0,
            description: "Test job".to_string(),
            final_loss: 0.0,
            best_validation_score: 0.0,
            tags: HashMap::new(),
        };

        state.add_job(job);
        state.update_job_status("test_job_2", TrainingStatus::Running);

        let updated = state.get_job("test_job_2").unwrap();
        assert_eq!(updated.status, TrainingStatus::Running as i32);
    }

    #[test]
    fn test_next_job_id() {
        let state = MockTrainingState::default();
        let id1 = state.next_job_id();
        let id2 = state.next_job_id();
        assert_ne!(id1, id2);
        assert_eq!(id1, "train_job_1");
        assert_eq!(id2, "train_job_2");
    }
}