Skip to main content

telos_agent/model/
mock.rs

1//! Mock provider for testing — dequeues pre-configured responses.
2//!
3//! [`MockProvider`] is constructed with a list of [`CompletionResponse`]s.
4//! Each call to [`complete`](MockProvider::complete) pops the next one from the
5//! queue and records the request it received. Running out of responses returns
6//! a provider error — making it easy to assert "the test exercised exactly N
7//! turns".
8
9use async_trait::async_trait;
10use std::collections::VecDeque;
11use tokio::sync::Mutex;
12
13use crate::error::AgentError;
14use crate::model::provider::{CompletionRequest, CompletionResponse, ModelProvider};
15
16/// A [`ModelProvider`] that returns pre-configured responses from a queue.
17pub struct MockProvider {
18    responses: Mutex<VecDeque<CompletionResponse>>,
19    /// All requests received, in arrival order. Public so tests can assert on them.
20    pub requests: Mutex<Vec<CompletionRequest>>,
21}
22
23impl MockProvider {
24    /// Build a mock that will reply with `responses` in FIFO order.
25    pub fn new(responses: Vec<CompletionResponse>) -> Self {
26        Self { responses: Mutex::new(responses.into()), requests: Mutex::new(Vec::new()) }
27    }
28}
29
30#[async_trait]
31impl ModelProvider for MockProvider {
32    async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, AgentError> {
33        self.requests.lock().await.push(request);
34        self.responses.lock().await.pop_front().ok_or_else(|| {
35            AgentError::Provider(crate::error::ProviderError::Other(
36                "mock provider has no more responses".into(),
37            ))
38        })
39    }
40}