vtcode_core/llm/
mock_client.rs

1//! Utilities for deterministic tests that need an `LLMClient` implementation
2//! without performing network calls.
3//!
4//! Enable the crate's `mock` feature to access the [`StaticResponseClient`],
5//! queue canned responses, and verify the interaction contract used by
6//! downstream integrations.
7
8use std::collections::VecDeque;
9
10use crate::llm::client::LLMClient;
11use async_trait::async_trait;
12use vtcode_commons::llm::{LLMError, LLMResponse};
13
14/// Deterministic `LLMClient` that yields queued responses.
15#[derive(Debug)]
16pub struct StaticResponseClient {
17    model: String,
18    queue: VecDeque<Result<LLMResponse, LLMError>>,
19}
20
21impl StaticResponseClient {
22    /// Create a mock client for the provided model.
23    pub fn new(model: impl Into<String>) -> Self {
24        Self { model: model.into(), queue: VecDeque::new() }
25    }
26
27    /// Queue a successful response. Responses are returned in FIFO order.
28    pub fn enqueue_response(&mut self, response: LLMResponse) {
29        self.queue.push_back(Ok(response));
30    }
31
32    /// Queue a successful response and return the client for chaining.
33    pub fn with_response(mut self, response: LLMResponse) -> Self {
34        self.enqueue_response(response);
35        self
36    }
37
38    /// Queue an error result. Errors are returned in FIFO order alongside responses.
39    pub fn enqueue_error(&mut self, error: LLMError) {
40        self.queue.push_back(Err(error));
41    }
42
43    /// Queue an error result and return the client for chaining.
44    pub fn with_error(mut self, error: LLMError) -> Self {
45        self.enqueue_error(error);
46        self
47    }
48
49    /// Consume the client and return it as a boxed trait object.
50    pub fn into_client(self) -> crate::llm::client::AnyClient {
51        Box::new(self)
52    }
53}
54
55#[async_trait]
56impl LLMClient for StaticResponseClient {
57    async fn generate(&mut self, _prompt: &str) -> Result<LLMResponse, LLMError> {
58        self.queue.pop_front().unwrap_or_else(|| {
59            Err(LLMError::InvalidRequest {
60                message: "StaticResponseClient has no queued responses".to_string(),
61                metadata: None,
62            })
63        })
64    }
65
66    fn model_id(&self) -> &str {
67        &self.model
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::StaticResponseClient;
74    use crate::llm::client::LLMClient;
75    use vtcode_commons::llm::{FinishReason, LLMError, LLMResponse};
76
77    #[test]
78    fn returns_responses_in_fifo_order() {
79        let response_one = LLMResponse {
80            content: Some("first".to_string()),
81            tool_calls: None,
82            model: "test".to_string(),
83            usage: None,
84            finish_reason: FinishReason::Stop,
85            reasoning: None,
86            reasoning_details: None,
87            tool_references: Vec::new(),
88            request_id: None,
89            organization_id: None,
90            compaction: None,
91        };
92        let response_two = LLMResponse {
93            content: Some("second".to_string()),
94            tool_calls: None,
95            model: "test".to_string(),
96            usage: None,
97            finish_reason: FinishReason::Stop,
98            reasoning: None,
99            reasoning_details: None,
100            tool_references: Vec::new(),
101            request_id: None,
102            organization_id: None,
103            compaction: None,
104        };
105
106        let mut client = StaticResponseClient::new("test");
107        client.enqueue_response(response_one.clone());
108        client.enqueue_response(response_two.clone());
109
110        let first = futures::executor::block_on(client.generate("prompt")).unwrap();
111        let second = futures::executor::block_on(client.generate("prompt")).unwrap();
112
113        assert_eq!(first.content, response_one.content);
114        assert_eq!(second.content, response_two.content);
115    }
116
117    #[test]
118    fn errors_when_queue_is_empty() {
119        let mut client = StaticResponseClient::new("test");
120        let error =
121            futures::executor::block_on(client.generate("prompt")).expect_err("expected error when queue is empty");
122
123        assert!(matches!(error, LLMError::InvalidRequest { .. }));
124    }
125}
vtcode_core/llm/mock_client.rs

vtcode_core/llm/
mock_client.rs
