//! Shared test utilities for zclaw-runtime and dependent crates. //! //! Provides `MockLlmDriver` — a controllable LLM driver for offline testing. use crate::driver::{ CompletionRequest, CompletionResponse, ContentBlock, LlmDriver, StopReason, }; use crate::stream::StreamChunk; use async_trait::async_trait; use futures::{Stream, StreamExt}; use serde_json::Value; use std::collections::VecDeque; use std::pin::Pin; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use zclaw_types::Result; use zclaw_types::ZclawError; /// Thread-safe mock LLM driver for testing. /// /// # Usage /// ```ignore /// let mock = MockLlmDriver::new() /// .with_text_response("Hello!") /// .with_text_response("How can I help?"); /// /// let resp = mock.complete(request).await?; /// assert_eq!(resp.content_text(), "Hello!"); /// ``` pub struct MockLlmDriver { responses: Arc>>, stream_chunks: Arc>>>, call_count: AtomicUsize, last_request: Arc>>, /// If true, `complete()` returns an error instead of a response. fail_mode: Arc>, } impl MockLlmDriver { pub fn new() -> Self { Self { responses: Arc::new(Mutex::new(VecDeque::new())), stream_chunks: Arc::new(Mutex::new(VecDeque::new())), call_count: AtomicUsize::new(0), last_request: Arc::new(Mutex::new(None)), fail_mode: Arc::new(Mutex::new(false)), } } /// Queue a text response. pub fn with_text_response(mut self, text: &str) -> Self { self.push_response(CompletionResponse { content: vec![ContentBlock::Text { text: text.to_string() }], model: "mock-model".to_string(), input_tokens: 10, output_tokens: text.len() as u32 / 4, stop_reason: StopReason::EndTurn, }); self } /// Queue a response with tool calls. pub fn with_tool_call(mut self, tool_name: &str, args: Value) -> Self { self.push_response(CompletionResponse { content: vec![ ContentBlock::Text { text: format!("Calling {}", tool_name) }, ContentBlock::ToolUse { id: format!("call_{}", self.call_count()), name: tool_name.to_string(), input: args, }, ], model: "mock-model".to_string(), input_tokens: 10, output_tokens: 20, stop_reason: StopReason::ToolUse, }); self } /// Queue an error response. pub fn with_error(mut self, _error: &str) -> Self { self.push_response(CompletionResponse { content: vec![], model: "mock-model".to_string(), input_tokens: 0, output_tokens: 0, stop_reason: StopReason::Error, }); self } /// Queue a raw response. pub fn with_response(mut self, response: CompletionResponse) -> Self { self.push_response(response); self } /// Queue stream chunks for a streaming call. pub fn with_stream_chunks(self, chunks: Vec) -> Self { self.stream_chunks .lock() .expect("stream_chunks lock") .push_back(chunks); self } /// Enable fail mode — all `complete()` calls return an error. pub fn set_fail_mode(&self, fail: bool) { *self.fail_mode.lock().expect("fail_mode lock") = fail; } /// Number of times `complete()` was called. pub fn call_count(&self) -> usize { self.call_count.load(Ordering::SeqCst) } /// Inspect the last request sent to the driver. pub fn last_request(&self) -> Option { self.last_request .lock() .expect("last_request lock") .clone() } fn push_response(&mut self, resp: CompletionResponse) { self.responses .lock() .expect("responses lock") .push_back(resp); } fn next_response(&self) -> CompletionResponse { let mut queue = self.responses.lock().expect("responses lock"); queue .pop_front() .unwrap_or_else(|| CompletionResponse { content: vec![ContentBlock::Text { text: "mock default response".to_string(), }], model: "mock-model".to_string(), input_tokens: 0, output_tokens: 0, stop_reason: StopReason::EndTurn, }) } } impl Default for MockLlmDriver { fn default() -> Self { Self::new() } } #[async_trait] impl LlmDriver for MockLlmDriver { fn provider(&self) -> &str { "mock" } async fn complete(&self, request: CompletionRequest) -> Result { self.call_count.fetch_add(1, Ordering::SeqCst); *self.last_request.lock().expect("last_request lock") = Some(request); if *self.fail_mode.lock().expect("fail_mode lock") { return Err(ZclawError::LlmError("mock driver fail mode".to_string())); } Ok(self.next_response()) } fn stream( &self, request: CompletionRequest, ) -> Pin> + Send + '_>> { self.call_count.fetch_add(1, Ordering::SeqCst); *self.last_request.lock().expect("last_request lock") = Some(request); let chunks: Vec> = self .stream_chunks .lock() .expect("stream_chunks lock") .pop_front() .unwrap_or_else(|| { vec![ StreamChunk::TextDelta { delta: "mock stream".to_string(), }, StreamChunk::Complete { input_tokens: 10, output_tokens: 2, stop_reason: "end_turn".to_string(), }, ] }) .into_iter() .map(Ok) .collect(); futures::stream::iter(chunks).boxed() } fn is_configured(&self) -> bool { true } }