use super::*;
use std::io::{Read, Write};
use std::net::TcpListener;
fn serve(content_type: &'static str, reply: String) -> (String, std::sync::mpsc::Receiver<String>) {
let listener = TcpListener::bind("127.0.0.1:0").expect("puerto libre");
let url = format!("http://{}/v1/fim/completions", listener.local_addr().expect("dirección"));
let (sender, receiver) = std::sync::mpsc::channel();
std::thread::spawn(move || {
let (mut socket, _) = listener.accept().expect("conexión");
let mut received = Vec::new();
let mut buffer = [0u8; 8192];
loop {
let read = socket.read(&mut buffer).expect("lectura");
received.extend_from_slice(&buffer[..read]);
let text = String::from_utf8_lossy(&received).to_string();
let Some(head) = text.find("\r\n\r\n") else { continue };
let length = text
.lines()
.find_map(|line| line.to_lowercase().strip_prefix("content-length:").map(|v| v.trim().to_owned()))
.and_then(|value| value.parse::<usize>().ok())
.unwrap_or(0);
if received.len() >= head + 4 + length {
let _ = sender.send(text[head + 4..].to_owned());
break;
}
}
let head = format!("HTTP/1.1 200 OK\r\ncontent-type: {content_type}\r\nconnection: close\r\n\r\n");
let _ = socket.write_all(format!("{head}{reply}").as_bytes());
});
(url, receiver)
}
fn collect(provider: &FimProvider) -> Result<String, AiError> {
let mut text = String::new();
provider.complete(
"fn suma(a: i32, b: i32) -> i32 {\n ",
"\n}\n",
64,
&CancelToken::default(),
&mut |event| {
if let Event::Text(chunk) = event {
text.push_str(&chunk);
}
},
)?;
Ok(text)
}
#[test]
fn the_three_backends_share_one_request_shape() {
let codestral = FimProvider::codestral("k".into());
let body = codestral.body("antes", "después", 64);
assert_eq!(body["model"], "codestral-latest");
assert_eq!((body["prompt"].as_str(), body["suffix"].as_str()), (Some("antes"), Some("después")));
assert_eq!((body["max_tokens"].as_u64(), body["stream"].as_bool()), (Some(64), Some(true)));
assert_eq!(codestral.url, "https://codestral.mistral.ai/v1/fim/completions");
assert!(!FimProvider::mercury("k".into()).stream, "Mercury no documenta streaming");
assert_eq!(FimProvider::ollama(None).url, "http://127.0.0.1:11434/v1/completions");
assert_eq!(FimProvider::ollama(Some("http://host:1/v1/")).url, "http://host:1/v1/completions");
assert_eq!(default_model("mercury"), Some("mercury-edit-2"));
assert_eq!(default_model("openai"), None);
}
#[test]
fn text_is_read_from_every_response_shape() {
assert_eq!(text_of(&json!({"choices": [{"text": "a + b"}]})), Some("a + b"));
assert_eq!(text_of(&json!({"choices": [{"delta": {"content": "a"}}]})), Some("a"));
assert_eq!(text_of(&json!({"choices": [{"message": {"content": "b"}}]})), Some("b"));
assert_eq!(text_of(&json!({"choices": []})), None);
}
#[test]
fn a_streamed_reply_arrives_in_pieces_and_the_key_is_sent() {
let chunk = |text: &str, done: bool| {
let finish = if done { json!("stop") } else { Value::Null };
format!("data: {}\n\n", json!({"choices": [{"index": 0, "delta": {"content": text}, "finish_reason": finish}]}))
};
let reply = format!("{}{}data: [DONE]\n\n", chunk("a ", false), chunk("+ b", true));
let (url, received) = serve("text/event-stream", reply);
let provider = FimProvider { url, ..FimProvider::codestral("clave-de-prueba".into()) };
assert_eq!(collect(&provider).unwrap(), "a + b");
let sent: Value = serde_json::from_str(&received.recv().unwrap()).unwrap();
assert_eq!(sent["suffix"], "\n}\n");
}
#[test]
fn a_whole_reply_without_streaming_also_works() {
let reply = json!({"choices": [{"index": 0, "text": "a + b", "finish_reason": "stop"}]}).to_string();
let (url, _received) = serve("application/json", reply);
let provider = FimProvider { url, ..FimProvider::mercury("k".into()) };
assert_eq!(collect(&provider).unwrap(), "a + b");
}
#[test]
fn an_error_inside_the_stream_is_reported() {
let mut decoder = FimDecoder::default();
let event = SseEvent { event: None, data: json!({"error": {"message": "clave no válida"}}).to_string() };
assert!(
matches!(decoder.feed(&event, &mut Vec::new()), Err(AiError::Api(message)) if message == "clave no válida")
);
}