#![allow(dead_code)] use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use async_trait::async_trait; use futures::Stream; use llm_worker::llm_client::event::Event; use llm_worker::llm_client::{ClientError, LlmClient, Request}; /// A mock LLM client that replays pre-defined event sequences. #[derive(Clone)] pub struct MockLlmClient { responses: Arc>>, call_count: Arc, } impl MockLlmClient { pub fn new(events: Vec) -> Self { Self::with_responses(vec![events]) } pub fn with_responses(responses: Vec>) -> Self { Self { responses: Arc::new(responses), call_count: Arc::new(AtomicUsize::new(0)), } } } #[async_trait] impl LlmClient for MockLlmClient { fn clone_boxed(&self) -> Box { Box::new(self.clone()) } async fn stream( &self, _request: Request, ) -> Result> + Send>>, ClientError> { let count = self.call_count.fetch_add(1, Ordering::SeqCst); if count >= self.responses.len() { return Err(ClientError::Api { status: Some(500), code: Some("mock_error".to_string()), message: "No more mock responses".to_string(), retry_after: None, }); } let events = self.responses[count].clone(); let stream = futures::stream::iter(events.into_iter().map(Ok)); Ok(Box::pin(stream)) } }