use crate::api::ApiClient;
use crate::api::error::ApiError;
use crate::config::LoopConfig;
use crate::message::{Message, MessagePart, Role};
use crate::stream::{
DeltaPart, IndexedDelta, MessageDelta, MessageDeltaPayload, MessageMetadata, MessageStart,
PartStart, StreamEvent, Usage,
};
use crate::tool::{Tool, ToolContext, ToolError, ToolOutput, ToolSchema};
use futures::Stream;
use serde_json::{Value, json};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use parking_lot::Mutex;
use uuid::Uuid;
#[derive(Clone)]
pub struct MockApiClient {
model_name: Arc<parking_lot::Mutex<String>>,
responses: Arc<Mutex<Vec<MockResponse>>>,
error: Option<String>,
}
#[derive(Clone, Default)]
pub struct MockResponse {
pub text: String,
pub tool_call: Option<MockToolCall>,
pub stop_reason: String,
}
#[derive(Clone)]
pub struct MockToolCall {
pub id: String,
pub name: String,
pub input: Value,
}
impl MockApiClient {
#[must_use]
pub fn new(model: &str) -> Self {
let default_response = MockResponse {
text: "Hello!".to_string(),
tool_call: None,
stop_reason: "end_turn".to_string(),
};
Self {
model_name: Arc::new(parking_lot::Mutex::new(model.to_string())),
responses: Arc::new(Mutex::new(vec![default_response])),
error: None,
}
}
#[must_use]
pub fn with_text_response(self, text: &str) -> Self {
if let Some(r) = self.responses.lock().first_mut() {
r.text = text.to_string();
}
self
}
#[must_use]
pub fn with_tool_call(self, id: &str, name: &str, input: Value) -> Self {
let mut responses = self.responses.lock();
if let Some(r) = responses.first_mut() {
r.tool_call = Some(MockToolCall {
id: id.to_string(),
name: name.to_string(),
input,
});
r.stop_reason = "tool_use".to_string();
}
drop(responses);
self
}
#[must_use]
pub fn with_stop_reason(self, reason: &str) -> Self {
if let Some(r) = self.responses.lock().first_mut() {
r.stop_reason = reason.to_string();
}
self
}
#[must_use]
pub fn with_responses(self, responses: Vec<MockResponse>) -> Self {
if !responses.is_empty() {
*self.responses.lock() = responses;
}
self
}
#[must_use]
pub fn with_error(mut self, error: &str) -> Self {
self.error = Some(error.to_string());
self
}
fn pop_response(&self) -> MockResponse {
let mut guard = self.responses.lock();
if guard.len() > 1 {
guard.remove(0)
} else {
guard.first().cloned().unwrap_or_default()
}
}
}
impl ApiClient for MockApiClient {
fn model(&self) -> String {
self.model_name.lock().clone()
}
fn set_model(&self, model: &str) -> bool {
if model.trim().is_empty() {
return false;
}
*self.model_name.lock() = model.to_string();
true
}
fn stream_messages(
&self,
_messages: Vec<Message>,
_system: Option<String>,
_tools: Option<Vec<ToolSchema>>,
) -> Pin<Box<dyn Stream<Item = Result<StreamEvent, ApiError>> + Send + 'static>> {
if let Some(ref err) = self.error {
let err = err.clone();
return Box::pin(futures::stream::once(
async move { Err(ApiError::api(&err)) },
));
}
let response = self.pop_response();
let model = self.model_name.lock().clone();
let mut events: Vec<Result<StreamEvent, ApiError>> =
vec![Ok(StreamEvent::MessageStart(MessageStart {
message: MessageMetadata {
id: "msg_test".to_string(),
role: "assistant".to_string(),
model,
},
}))];
let text = response.text.clone();
events.push(Ok(StreamEvent::PartStart(PartStart {
index: 0,
part: Some(MessagePart::text("")),
})));
events.push(Ok(StreamEvent::IndexedDelta(IndexedDelta {
index: 0,
delta: DeltaPart::Text { text },
})));
events.push(Ok(StreamEvent::PartStop));
if let Some(tc) = &response.tool_call {
events.push(Ok(StreamEvent::PartStart(PartStart {
index: 1,
part: Some(MessagePart::tool_call(&tc.id, &tc.name, tc.input.clone())),
})));
events.push(Ok(StreamEvent::PartStop));
}
events.push(Ok(StreamEvent::MessageDelta(MessageDelta {
delta: MessageDeltaPayload {
stop_reason: Some(response.stop_reason),
},
usage: Some(Usage::new(50, 25)),
})));
events.push(Ok(StreamEvent::MessageStop));
Box::pin(futures::stream::iter(events))
}
fn create_message(
&self,
_messages: Vec<Message>,
_system: Option<String>,
_tools: Option<Vec<ToolSchema>>,
) -> Pin<Box<dyn Future<Output = Result<Value, ApiError>> + Send + '_>> {
if let Some(ref err) = self.error {
let err = err.clone();
return Box::pin(async move { Err(ApiError::api(&err)) });
}
let response = self.pop_response();
Box::pin(async move {
Ok(json!({
"content": [{"type": "text", "text": response.text}]
}))
})
}
}
pub struct MockTool {
name: String,
description: String,
input_schema: Value,
result: String,
is_error: bool,
is_concurrency_safe: bool,
is_read_only: bool,
system_prompt: Option<String>,
}
impl MockTool {
#[must_use]
pub fn new(name: &str, description: &str) -> Self {
Self {
name: name.to_string(),
description: description.to_string(),
input_schema: json!({
"type": "object",
"properties": { "input": { "type": "string" } }
}),
result: "mock result".to_string(),
is_error: false,
is_concurrency_safe: false,
is_read_only: true,
system_prompt: None,
}
}
#[must_use]
pub fn with_result(mut self, result: &str) -> Self {
self.result = result.to_string();
self
}
#[must_use]
pub fn with_error(mut self) -> Self {
self.is_error = true;
self
}
#[must_use]
pub fn with_concurrency_safe(mut self, safe: bool) -> Self {
self.is_concurrency_safe = safe;
self
}
#[must_use]
pub fn with_read_only(mut self, read_only: bool) -> Self {
self.is_read_only = read_only;
self
}
#[must_use]
pub fn with_schema(mut self, schema: Value) -> Self {
self.input_schema = schema;
self
}
#[must_use]
pub fn with_system_prompt(mut self, prompt: &str) -> Self {
self.system_prompt = Some(prompt.to_string());
self
}
}
impl Tool for MockTool {
fn name(&self) -> &str {
&self.name
}
fn description(&self) -> &str {
&self.description
}
fn schema(&self) -> ToolSchema {
ToolSchema {
tool: self.name.clone(),
description: self.description.clone(),
input_schema: self.input_schema.clone(),
}
}
fn call(
&self,
_input: Value,
_context: &ToolContext,
) -> Pin<Box<dyn Future<Output = Result<ToolOutput, ToolError>> + Send + '_>> {
let result = self.result.clone();
let is_error = self.is_error;
Box::pin(async move {
if is_error {
Err(ToolError::Execution(result))
} else {
Ok(ToolOutput::text(result))
}
})
}
fn is_concurrency_safe(&self) -> bool {
self.is_concurrency_safe
}
fn is_read_only(&self) -> bool {
self.is_read_only
}
fn system_prompt(&self) -> Option<String> {
self.system_prompt.clone()
}
}
#[must_use]
pub fn test_message(text: &str) -> Message {
Message::user(text)
}
#[must_use]
pub fn test_assistant_message(text: &str) -> Message {
Message::assistant(text)
}
#[must_use]
pub fn test_tool_use_message(calls: &[(&str, &str, Value)]) -> Message {
let blocks: Vec<MessagePart> = calls
.iter()
.enumerate()
.map(|(i, (id, name, input))| {
MessagePart::tool_call(
if id.is_empty() {
format!("call_{i}")
} else {
id.to_string()
},
*name,
input.clone(),
)
})
.collect();
Message::new(Role::Assistant, blocks)
}
#[must_use]
pub fn test_config() -> LoopConfig {
LoopConfig {
session_id: Uuid::new_v4(),
max_turns: 10,
system_prompt: Some("You are a test assistant.".to_string()),
..Default::default()
}
}
#[must_use]
pub fn test_config_with_id(id: Uuid) -> LoopConfig {
LoopConfig {
session_id: id,
max_turns: 10,
system_prompt: Some("You are a test assistant.".to_string()),
..Default::default()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::tool::ToolRegistry;
use futures::StreamExt;
#[test]
fn test_mock_client_model() {
let client = MockApiClient::new("test-model");
assert_eq!(client.model(), "test-model");
}
#[tokio::test]
async fn test_mock_client_default_response() {
let client = MockApiClient::new("test-model");
let stream = client.stream_messages(vec![Message::user("Hi")], None, None);
let events: Vec<_> = stream.collect().await;
assert!(events.len() >= 4);
assert!(events[0].is_ok());
}
#[tokio::test]
async fn test_mock_client_custom_text() {
let client = MockApiClient::new("test-model").with_text_response("Custom response");
let stream = client.stream_messages(vec![Message::user("Hi")], None, None);
let events: Vec<_> = stream.collect().await;
let has_text = events.iter().any(|e| {
if let Ok(StreamEvent::IndexedDelta(delta)) = e {
if let DeltaPart::Text { text } = &delta.delta {
return text == "Custom response";
}
}
false
});
assert!(has_text);
}
#[tokio::test]
async fn test_mock_client_tool_call() {
let client = MockApiClient::new("test-model").with_tool_call(
"call_1",
"echo",
json!({"message": "hi"}),
);
let stream = client.stream_messages(vec![Message::user("Hi")], None, None);
let events: Vec<_> = stream.collect().await;
let has_tool_use = events.iter().any(|e| {
if let Ok(StreamEvent::PartStart(start)) = e {
if let Some(MessagePart::ToolCall { name, .. }) = &start.part {
return name == "echo";
}
}
false
});
assert!(has_tool_use);
let has_tool_stop = events.iter().any(|e| {
if let Ok(StreamEvent::MessageDelta(delta)) = e {
delta.delta.stop_reason.as_deref() == Some("tool_use")
} else {
false
}
});
assert!(has_tool_stop);
}
#[tokio::test]
async fn test_mock_client_error() {
let client = MockApiClient::new("test-model").with_error("API error");
let stream = client.stream_messages(vec![Message::user("Hi")], None, None);
let events: Vec<_> = stream.collect().await;
assert_eq!(events.len(), 1);
assert!(events[0].is_err());
}
#[tokio::test]
async fn test_mock_client_multi_turn() {
let client = MockApiClient::new("test-model").with_responses(vec![
MockResponse {
text: "First".to_string(),
tool_call: None,
stop_reason: "end_turn".to_string(),
},
MockResponse {
text: "Second".to_string(),
tool_call: None,
stop_reason: "end_turn".to_string(),
},
]);
let stream1 = client.stream_messages(vec![Message::user("Hi")], None, None);
let events1: Vec<_> = stream1.collect().await;
let has_first = events1.iter().any(|e| {
if let Ok(StreamEvent::IndexedDelta(delta)) = e {
if let DeltaPart::Text { text } = &delta.delta {
return text == "First";
}
}
false
});
assert!(has_first);
let stream2 = client.stream_messages(vec![Message::user("Hi")], None, None);
let events2: Vec<_> = stream2.collect().await;
let has_second = events2.iter().any(|e| {
if let Ok(StreamEvent::IndexedDelta(delta)) = e {
if let DeltaPart::Text { text } = &delta.delta {
return text == "Second";
}
}
false
});
assert!(has_second);
}
#[tokio::test]
async fn test_mock_client_create_message() {
let client = MockApiClient::new("test-model").with_text_response("Hello!");
let result = client
.create_message(vec![Message::user("Hi")], None, None)
.await;
assert!(result.is_ok());
let json = result.unwrap();
assert_eq!(json["content"][0]["text"], "Hello!");
}
#[tokio::test]
async fn test_mock_tool() {
let tool = MockTool::new("echo", "Echoes input").with_result("Echo: hello");
let ctx = ToolContext::default();
let result = tool.call(json!({"input": "hello"}), &ctx).await;
assert!(result.is_ok());
assert_eq!(result.unwrap().text_content(), "Echo: hello");
}
#[tokio::test]
async fn test_mock_tool_error() {
let tool = MockTool::new("fail", "Always fails")
.with_result("Something went wrong")
.with_error();
let ctx = ToolContext::default();
let result = tool.call(json!({}), &ctx).await;
assert!(result.is_err());
}
#[test]
fn test_mock_tool_registry() {
let tool = MockTool::new("echo", "Echoes input").with_concurrency_safe(true);
let mut registry = ToolRegistry::new();
registry.register(tool);
assert!(registry.contains("echo"));
assert_eq!(registry.len(), 1);
let t = registry.get("echo").unwrap();
assert_eq!(t.name(), "echo");
assert!(t.is_concurrency_safe());
assert!(t.is_read_only());
}
#[test]
fn test_fixture_test_message() {
let msg = test_message("Hello");
assert_eq!(msg.role, Role::User);
assert_eq!(msg.parts.len(), 1);
}
#[test]
fn test_fixture_test_assistant_message() {
let msg = test_assistant_message("Hi there");
assert_eq!(msg.role, Role::Assistant);
}
#[test]
fn test_fixture_test_tool_use_message() {
let msg = test_tool_use_message(&[("call_1", "bash", json!({"command": "ls"}))]);
assert_eq!(msg.role, Role::Assistant);
assert_eq!(msg.parts.len(), 1);
assert!(msg.parts[0].is_tool_call());
}
#[test]
fn test_fixture_test_config() {
let config = test_config();
assert_eq!(config.max_turns, 10);
assert!(config.system_prompt.is_some());
}
#[test]
fn mock_api_client_set_model() {
let client = MockApiClient::new("model-a");
assert_eq!(client.model(), "model-a");
assert!(client.set_model("model-b"));
assert_eq!(client.model(), "model-b");
}
#[test]
fn mock_api_client_set_model_rejects_empty() {
let client = MockApiClient::new("model-a");
assert!(!client.set_model(""));
assert!(!client.set_model(" "));
assert_eq!(client.model(), "model-a");
}
}