1use std::sync::Mutex;
8
9use async_trait::async_trait;
10
11use super::{Completion, LlmProvider, ToolSpec};
12use crate::error::{Error, Result};
13use crate::message::Message;
14
15#[derive(Default)]
16pub struct MockProvider {
17 scripted: Mutex<Vec<Completion>>,
18 calls: Mutex<Vec<Vec<Message>>>,
19}
20
21impl MockProvider {
22 pub fn new(scripted: Vec<Completion>) -> Self {
23 Self {
24 scripted: Mutex::new(scripted),
25 calls: Mutex::new(Vec::new()),
26 }
27 }
28
29 pub fn calls(&self) -> Vec<Vec<Message>> {
31 self.calls.lock().unwrap().clone()
32 }
33}
34
35#[async_trait]
36impl LlmProvider for MockProvider {
37 async fn complete(&self, messages: &[Message], _tools: &[ToolSpec]) -> Result<Completion> {
38 self.calls.lock().unwrap().push(messages.to_vec());
39 let mut queue = self.scripted.lock().unwrap();
40 if queue.is_empty() {
41 return Err(Error::Llm(
42 "MockProvider: no scripted completions left".into(),
43 ));
44 }
45 Ok(queue.remove(0))
46 }
47}
48
49#[cfg(test)]
50mod tests {
51 use super::*;
52
53 #[tokio::test]
54 async fn returns_scripted_in_order_and_records_calls() {
55 let provider = MockProvider::new(vec![
56 Completion {
57 content: "one".into(),
58 tool_calls: vec![],
59 finish_reason: Some("stop".into()),
60 },
61 Completion {
62 content: "two".into(),
63 tool_calls: vec![],
64 finish_reason: Some("stop".into()),
65 },
66 ]);
67
68 let m1 = vec![Message::user("hi")];
69 let r1 = provider.complete(&m1, &[]).await.unwrap();
70 assert_eq!(r1.content, "one");
71
72 let m2 = vec![Message::user("bye")];
73 let r2 = provider.complete(&m2, &[]).await.unwrap();
74 assert_eq!(r2.content, "two");
75
76 assert_eq!(provider.calls().len(), 2);
77 assert_eq!(provider.calls()[0][0].content, "hi");
78 }
79
80 #[tokio::test]
81 async fn errors_when_queue_drained() {
82 let provider = MockProvider::new(vec![]);
83 let err = provider.complete(&[], &[]).await.unwrap_err();
84 assert!(matches!(err, Error::Llm(_)));
85 }
86}