test: add Phase 2 streaming SSE passthrough tests
Confirms the existing proxy streams SSE chunks incrementally: - 5-chunk test with 50ms delays verifies time spread between first and last chunk arrival (not buffered) - Verifies data: [DONE] terminator is forwarded No src/ changes needed — Body::from_stream(bytes_stream()) already handles SSE correctly. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -1,3 +1,8 @@
|
||||
#![allow(dead_code)]
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::http::header;
|
||||
use axum::response::Response;
|
||||
use axum::routing::{get, post};
|
||||
use axum::{Json, Router};
|
||||
use cortex_core::config::{
|
||||
@@ -5,8 +10,10 @@ use cortex_core::config::{
|
||||
};
|
||||
use cortex_core::node::{ModelEntry, ModelStatus};
|
||||
use cortex_gateway::state::CortexState;
|
||||
use futures::{StreamExt, stream};
|
||||
use serde_json::{Value, json};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
/// Spawns a mock mistral.rs backend on a random port.
|
||||
@@ -63,6 +70,66 @@ async fn mock_list_models() -> Json<Value> {
|
||||
}))
|
||||
}
|
||||
|
||||
/// Spawns a mock mistral.rs backend that returns SSE streaming responses.
|
||||
/// Each chunk is delayed by `chunk_delay` to prove the proxy streams incrementally.
|
||||
pub async fn spawn_streaming_mock_backend(chunk_count: usize, chunk_delay: Duration) -> String {
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
post(move |Json(body): Json<Value>| async move {
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("unknown")
|
||||
.to_string();
|
||||
|
||||
let chunks: Vec<String> = (0..chunk_count)
|
||||
.map(|i| {
|
||||
let content = format!("token{i}");
|
||||
let chunk = json!({
|
||||
"id": "chatcmpl-stream-001",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1700000000_u64,
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": { "content": content },
|
||||
"finish_reason": null
|
||||
}]
|
||||
});
|
||||
format!("data: {chunk}\n\n")
|
||||
})
|
||||
.collect();
|
||||
|
||||
let delay = chunk_delay;
|
||||
let stream = stream::iter(
|
||||
chunks
|
||||
.into_iter()
|
||||
.chain(std::iter::once("data: [DONE]\n\n".to_string())),
|
||||
)
|
||||
.then(move |chunk| async move {
|
||||
tokio::time::sleep(delay).await;
|
||||
Ok::<_, std::convert::Infallible>(chunk)
|
||||
});
|
||||
|
||||
Response::builder()
|
||||
.header(header::CONTENT_TYPE, "text/event-stream")
|
||||
.header(header::CACHE_CONTROL, "no-cache")
|
||||
.body(Body::from_stream(stream))
|
||||
.unwrap()
|
||||
}),
|
||||
)
|
||||
.route("/v1/models", get(mock_list_models));
|
||||
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
|
||||
format!("http://{addr}")
|
||||
}
|
||||
|
||||
/// Spawns the cortex gateway with a single node pointing at `mock_url`.
|
||||
/// The node is pre-seeded as healthy with one loaded model ("test-model").
|
||||
/// Returns the gateway's base URL.
|
||||
|
||||
Reference in New Issue
Block a user