use std::collections::VecDeque;
use std::sync::{Mutex, PoisonError};
use async_trait::async_trait;
use crate::completion::{Completion, CompletionDelta};
use crate::provider::{Provider, ProviderError};
use crate::request::ConversationRequest;
pub struct MockProvider {
script: Mutex<VecDeque<Result<Completion, ProviderError>>>,
}
impl MockProvider {
#[must_use]
pub fn new(script: Vec<Completion>) -> Self {
Self {
script: Mutex::new(script.into_iter().map(Ok).collect()),
}
}
#[must_use]
pub fn with_results(script: Vec<Result<Completion, ProviderError>>) -> Self {
Self {
script: Mutex::new(script.into_iter().collect()),
}
}
}
#[async_trait]
impl Provider for MockProvider {
#[allow(clippy::unnecessary_literal_bound)]
fn api_schema(&self) -> &str {
"mock"
}
async fn complete(&self, _request: &ConversationRequest) -> Result<Completion, ProviderError> {
let mut script = self.script.lock().unwrap_or_else(PoisonError::into_inner);
match script.pop_front() {
Some(result) => result,
None => panic!(
"MockProvider script exhausted: complete() was called more times than scripted"
),
}
}
async fn stream(
&self,
request: &ConversationRequest,
on_delta: &mut (dyn FnMut(CompletionDelta) + Send),
) -> Result<Completion, ProviderError> {
let completion = self.complete(request).await?;
if let Some(text) = completion.text() {
for chunk in chunk_text(&text) {
on_delta(CompletionDelta::Text(chunk));
}
}
Ok(completion)
}
}
fn chunk_text(text: &str) -> Vec<String> {
let mut chunks = Vec::new();
let mut cur = String::new();
let mut in_word = false;
for ch in text.chars() {
if ch.is_whitespace() {
cur.push(ch);
in_word = false;
} else {
if !in_word && !cur.is_empty() {
chunks.push(std::mem::take(&mut cur));
}
cur.push(ch);
in_word = true;
}
}
if !cur.is_empty() {
chunks.push(cur);
}
chunks
}
#[cfg(test)]
mod tests {
use super::*;
use crate::request::{CacheHint, SamplingArgs};
use locode_protocol::{ContentBlock, Usage};
use crate::completion::StopReason;
fn req() -> ConversationRequest {
ConversationRequest {
messages: vec![],
tools: vec![],
sampling_args: SamplingArgs::default(),
cache_hint: CacheHint::default(),
}
}
fn text(t: &str) -> Completion {
Completion {
content: vec![ContentBlock::Text { text: t.into() }],
usage: Usage::default(),
stop: StopReason::EndTurn,
}
}
#[test]
fn chunk_text_concatenates_back_to_input_exactly() {
for input in [
"",
"hello",
"hello world",
" leading and double spaces ",
"line one\nline two\n",
"tabs\tand\nnewlines here",
"unicode → café ☕ done",
] {
let joined: String = chunk_text(input).concat();
assert_eq!(joined, input, "chunk_text must reconstruct {input:?}");
}
}
#[test]
fn chunk_text_splits_on_word_boundaries() {
assert_eq!(chunk_text("all done"), vec!["all ", "done"]);
assert_eq!(chunk_text("one two three"), vec!["one ", "two ", "three"]);
assert_eq!(chunk_text("solo"), vec!["solo"]);
assert!(chunk_text("").is_empty());
}
#[tokio::test]
async fn mock_stream_chunks_text_and_returns_the_whole_completion() {
let mock = MockProvider::new(vec![text("hello there world")]);
let mut deltas = Vec::new();
let completion = mock
.stream(&req(), &mut |d| deltas.push(d))
.await
.expect("mock stream ok");
assert!(deltas.len() > 1, "expected multiple chunks: {deltas:?}");
let joined: String = deltas
.iter()
.map(|d| match d {
CompletionDelta::Text(t) => t.as_str(),
other => panic!("mock should only emit Text: {other:?}"),
})
.collect();
assert_eq!(joined, "hello there world");
assert_eq!(completion.text().as_deref(), Some("hello there world"));
assert_eq!(completion.stop, StopReason::EndTurn);
}
#[tokio::test]
async fn mock_stream_with_no_text_emits_no_deltas() {
let tool = Completion {
content: vec![ContentBlock::ToolUse {
id: "c1".into(),
name: "echo".into(),
input: serde_json::json!({}),
}],
usage: Usage::default(),
stop: StopReason::ToolUse,
};
let mock = MockProvider::new(vec![tool]);
let mut deltas = Vec::new();
let completion = mock
.stream(&req(), &mut |d| deltas.push(d))
.await
.expect("ok");
assert!(
deltas.is_empty(),
"tool-only turn streams no text: {deltas:?}"
);
assert!(completion.has_tool_calls());
}
}