adk_model/
mock.rs

1use adk_core::{Llm, LlmRequest, LlmResponse, LlmResponseStream, Result};
2use async_trait::async_trait;
3
4pub struct MockLlm {
5    name: String,
6    responses: Vec<LlmResponse>,
7}
8
9impl MockLlm {
10    pub fn new(name: impl Into<String>) -> Self {
11        Self { name: name.into(), responses: vec![] }
12    }
13
14    pub fn with_response(mut self, response: LlmResponse) -> Self {
15        self.responses.push(response);
16        self
17    }
18}
19
20#[async_trait]
21impl Llm for MockLlm {
22    fn name(&self) -> &str {
23        &self.name
24    }
25
26    async fn generate_content(&self, _req: LlmRequest, _stream: bool) -> Result<LlmResponseStream> {
27        let responses = self.responses.clone();
28        let stream = async_stream::stream! {
29            for response in responses {
30                yield Ok(response);
31            }
32        };
33        Ok(Box::pin(stream))
34    }
35}
36
37#[cfg(test)]
38mod tests {
39    use super::*;
40    use adk_core::Content;
41
42    #[test]
43    fn test_mock_llm() {
44        let mock =
45            MockLlm::new("test-llm").with_response(LlmResponse::new(Content::new("assistant")));
46        assert_eq!(mock.name(), "test-llm");
47        assert_eq!(mock.responses.len(), 1);
48    }
49
50    #[tokio::test]
51    async fn test_mock_llm_generate() {
52        use futures::StreamExt;
53
54        let mock = MockLlm::new("test")
55            .with_response(LlmResponse::new(Content::new("assistant").with_text("Hello")));
56
57        let req = LlmRequest::new("test", vec![]);
58        let mut stream = mock.generate_content(req, false).await.unwrap();
59
60        let response = stream.next().await.unwrap().unwrap();
61        assert!(response.content.is_some());
62    }
63}
adk_model/mock.rs

adk_model/
mock.rs
