gahow/agentic-pd-hybrid
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: vendor sglang v0.5.10 snapshot

Browse Source
This commit is contained in:
Gahow Wang
2026-04-24 12:29:36 +00:00
parent 78f0d15221
commit bded08301f
4308 changed files with 1200894 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

180
third_party/sglang/sgl-model-gateway/tests/common/mock_mcp_server.rs vendored Normal file
View File

@@ -0,0 +1,180 @@
// tests/common/mock_mcp_server.rs - Mock MCP server for testing
use rmcp::{
handler::server::{router::tool::ToolRouter, wrapper::Parameters},
model::*,
service::RequestContext,
tool, tool_handler, tool_router,
transport::streamable_http_server::{
session::local::LocalSessionManager, StreamableHttpService,
},
ErrorData as McpError, RoleServer, ServerHandler,
};
use tokio::net::TcpListener;
/// Mock MCP server that returns hardcoded responses for testing
pub struct MockMCPServer {
pub port: u16,
pub server_handle: Option<tokio::task::JoinHandle<()>>,
}
/// Simple test server with mock search tools
#[derive(Clone)]
pub struct MockSearchServer {
tool_router: ToolRouter<MockSearchServer>,
}
impl Default for MockSearchServer {
fn default() -> Self {
Self::new()
}
}
#[tool_router]
impl MockSearchServer {
pub fn new() -> Self {
Self {
tool_router: Self::tool_router(),
}
}
#[tool(description = "Mock web search tool")]
fn brave_web_search(
&self,
Parameters(params): Parameters<serde_json::Map<String, serde_json::Value>>,
) -> Result<CallToolResult, McpError> {
let query = params
.get("query")
.and_then(|v| v.as_str())
.unwrap_or("test");
Ok(CallToolResult::success(vec![Content::text(format!(
"Mock search results for: {}",
query
))]))
}
#[tool(description = "Mock local search tool")]
fn brave_local_search(
&self,
Parameters(_params): Parameters<serde_json::Map<String, serde_json::Value>>,
) -> Result<CallToolResult, McpError> {
Ok(CallToolResult::success(vec![Content::text(
"Mock local search results",
)]))
}
}
#[tool_handler]
impl ServerHandler for MockSearchServer {
fn get_info(&self) -> ServerInfo {
ServerInfo {
protocol_version: ProtocolVersion::V_2024_11_05,
capabilities: ServerCapabilities::builder().enable_tools().build(),
server_info: Implementation::from_build_env(),
instructions: Some("Mock server for testing".to_string()),
}
}
async fn initialize(
&self,
_request: InitializeRequestParam,
_context: RequestContext<RoleServer>,
) -> Result<InitializeResult, McpError> {
Ok(self.get_info())
}
}
impl MockMCPServer {
/// Start a mock MCP server on an available port
pub async fn start() -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
// Find an available port
let listener = TcpListener::bind("127.0.0.1:0").await?;
let port = listener.local_addr()?.port();
// Create the MCP service using rmcp's StreamableHttpService
let service = StreamableHttpService::new(
|| Ok(MockSearchServer::new()),
LocalSessionManager::default().into(),
Default::default(),
);
let app = axum::Router::new().nest_service("/mcp", service);
let server_handle = tokio::spawn(async move {
axum::serve(listener, app)
.await
.expect("Mock MCP server failed to start");
});
// Give the server a moment to start
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
Ok(MockMCPServer {
port,
server_handle: Some(server_handle),
})
}
/// Get the full URL for this mock server
pub fn url(&self) -> String {
format!("http://127.0.0.1:{}/mcp", self.port)
}
/// Stop the mock server
pub async fn stop(&mut self) {
if let Some(handle) = self.server_handle.take() {
handle.abort();
// Wait a moment for cleanup
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
}
}
}
impl Drop for MockMCPServer {
fn drop(&mut self) {
if let Some(handle) = self.server_handle.take() {
handle.abort();
}
}
}
#[cfg(test)]
mod tests {
#[allow(unused_imports)]
use super::MockMCPServer;
#[tokio::test]
async fn test_mock_server_startup() {
let mut server = MockMCPServer::start().await.unwrap();
assert!(server.port > 0);
assert!(server.url().contains(&server.port.to_string()));
server.stop().await;
}
#[tokio::test]
async fn test_mock_server_with_rmcp_client() {
let mut server = MockMCPServer::start().await.unwrap();
use rmcp::{transport::StreamableHttpClientTransport, ServiceExt};
let transport = StreamableHttpClientTransport::from_uri(server.url().as_str());
let client = ().serve(transport).await;
assert!(client.is_ok(), "Should be able to connect to mock server");
if let Ok(client) = client {
let tools = client.peer().list_all_tools().await;
assert!(tools.is_ok(), "Should be able to list tools");
if let Ok(tools) = tools {
assert_eq!(tools.len(), 2, "Should have 2 tools");
assert!(tools.iter().any(|t| t.name == "brave_web_search"));
assert!(tools.iter().any(|t| t.name == "brave_local_search"));
}
// Shutdown by dropping the client
drop(client);
}
server.stop().await;
}
}

240
third_party/sglang/sgl-model-gateway/tests/common/mock_openai_server.rs vendored Normal file
View File

@@ -0,0 +1,240 @@
//! Mock servers for testing
#![allow(dead_code)]
use std::{net::SocketAddr, sync::Arc};
use axum::{
body::Body,
extract::{Request, State},
http::{HeaderValue, StatusCode},
response::{
sse::{Event, KeepAlive},
IntoResponse, Response, Sse,
},
routing::post,
Json, Router,
};
use futures_util::stream::{self, StreamExt};
use serde_json::json;
use tokio::net::TcpListener;
/// Mock OpenAI API server for testing
pub struct MockOpenAIServer {
addr: SocketAddr,
_handle: tokio::task::JoinHandle<()>,
}
#[derive(Clone)]
struct MockServerState {
require_auth: bool,
expected_auth: Option<String>,
}
impl MockOpenAIServer {
/// Create and start a new mock OpenAI server
pub async fn new() -> Self {
Self::new_with_auth(None).await
}
/// Create and start a new mock OpenAI server with optional auth requirement
pub async fn new_with_auth(expected_auth: Option<String>) -> Self {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let state = Arc::new(MockServerState {
require_auth: expected_auth.is_some(),
expected_auth,
});
let app = Router::new()
.route("/v1/chat/completions", post(mock_chat_completions))
.route("/v1/completions", post(mock_completions))
.route("/v1/models", post(mock_models).get(mock_models))
.with_state(state);
let handle = tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
// Give the server a moment to start
tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
Self {
addr,
_handle: handle,
}
}
/// Get the base URL for this mock server
pub fn base_url(&self) -> String {
format!("http://{}", self.addr)
}
}
/// Mock chat completions endpoint
async fn mock_chat_completions(req: Request<Body>) -> Response {
let (_, body) = req.into_parts();
let body_bytes = match axum::body::to_bytes(body, usize::MAX).await {
Ok(bytes) => bytes,
Err(_) => return StatusCode::BAD_REQUEST.into_response(),
};
let request: serde_json::Value = match serde_json::from_slice(&body_bytes) {
Ok(req) => req,
Err(_) => return StatusCode::BAD_REQUEST.into_response(),
};
// Extract model from request or use default (owned String to satisfy 'static in stream)
let model: String = request
.get("model")
.and_then(|v| v.as_str())
.unwrap_or("gpt-3.5-turbo")
.to_string();
// If stream requested, return SSE
let is_stream = request
.get("stream")
.and_then(|v| v.as_bool())
.unwrap_or(false);
if is_stream {
let created = 1677652288u64;
// Single chunk then [DONE]
let model_chunk = model.clone();
let event_stream = stream::once(async move {
let chunk = json!({
"id": "chatcmpl-123456789",
"object": "chat.completion.chunk",
"created": created,
"model": model_chunk,
"choices": [{
"index": 0,
"delta": {
"content": "Hello!"
},
"finish_reason": null
}]
});
Ok::<_, std::convert::Infallible>(Event::default().data(chunk.to_string()))
})
.chain(stream::once(async { Ok(Event::default().data("[DONE]")) }));
Sse::new(event_stream)
.keep_alive(KeepAlive::default())
.into_response()
} else {
// Create a mock non-streaming response
let response = json!({
"id": "chatcmpl-123456789",
"object": "chat.completion",
"created": 1677652288,
"model": model,
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": "Hello! I'm a mock OpenAI assistant. How can I help you today?"
},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 9,
"completion_tokens": 12,
"total_tokens": 21
}
});
Json(response).into_response()
}
}
/// Mock completions endpoint (legacy)
async fn mock_completions(req: Request<Body>) -> Response {
let (_, body) = req.into_parts();
let body_bytes = match axum::body::to_bytes(body, usize::MAX).await {
Ok(bytes) => bytes,
Err(_) => return StatusCode::BAD_REQUEST.into_response(),
};
let request: serde_json::Value = match serde_json::from_slice(&body_bytes) {
Ok(req) => req,
Err(_) => return StatusCode::BAD_REQUEST.into_response(),
};
let model = request["model"].as_str().unwrap_or("text-davinci-003");
let response = json!({
"id": "cmpl-123456789",
"object": "text_completion",
"created": 1677652288,
"model": model,
"choices": [{
"text": " This is a mock completion response.",
"index": 0,
"logprobs": null,
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 5,
"completion_tokens": 7,
"total_tokens": 12
}
});
Json(response).into_response()
}
/// Mock models endpoint
async fn mock_models(State(state): State<Arc<MockServerState>>, req: Request<Body>) -> Response {
// Optionally enforce Authorization header
if state.require_auth {
let auth = req
.headers()
.get("authorization")
.or_else(|| req.headers().get("Authorization"))
.and_then(|v| v.to_str().ok())
.map(|s| s.to_string());
let auth_ok = match (&state.expected_auth, auth) {
(Some(expected), Some(got)) => &got == expected,
(None, Some(_)) => true,
_ => false,
};
if !auth_ok {
let mut response = Response::new(Body::from(
json!({
"error": {
"message": "Unauthorized",
"type": "invalid_request_error"
}
})
.to_string(),
));
*response.status_mut() = StatusCode::UNAUTHORIZED;
response
.headers_mut()
.insert("WWW-Authenticate", HeaderValue::from_static("Bearer"));
return response;
}
}
let response = json!({
"object": "list",
"data": [
{
"id": "gpt-4",
"object": "model",
"created": 1677610602,
"owned_by": "openai"
},
{
"id": "gpt-3.5-turbo",
"object": "model",
"created": 1677610602,
"owned_by": "openai"
}
]
});
Json(response).into_response()
}

1278
third_party/sglang/sgl-model-gateway/tests/common/mock_worker.rs vendored Executable file
View File

File diff suppressed because it is too large Load Diff

1028
third_party/sglang/sgl-model-gateway/tests/common/mod.rs vendored Normal file
View File

File diff suppressed because it is too large Load Diff

118
third_party/sglang/sgl-model-gateway/tests/common/redis_test_server.rs vendored Normal file
View File

@@ -0,0 +1,118 @@
use std::{
process::{Child, Command},
sync::OnceLock,
time::Duration,
};
use redis::RedisError;
use tracing::{info, warn};
static SHARED_SERVER: OnceLock<RedisTestServer> = OnceLock::new();
pub fn get_shared_server() -> &'static RedisTestServer {
let server = SHARED_SERVER
.get_or_init(|| RedisTestServer::start().expect("Failed to start shared Redis server"));
server.wait_ready();
server
}
pub struct RedisTestServer {
process: Option<Child>,
port: u16,
url: String,
}
impl RedisTestServer {
pub fn start() -> Result<Self, String> {
let port = portpicker::pick_unused_port()
.ok_or_else(|| "Failed to find available port".to_string())?;
let url = format!("redis://127.0.0.1:{}", port);
let mut cmd = Command::new("redis-server");
cmd.args([
"--port",
&port.to_string(),
"--save",
"",
"--appendonly",
"no",
"--daemonize",
"no",
]);
info!("Starting redis server... cmd={cmd:?}");
let process = cmd.spawn().map_err(|e| {
format!(
"Failed to start redis-server: {}. Is redis-server installed?",
e
)
})?;
Ok(Self {
process: Some(process),
port,
url,
})
}
pub fn wait_ready(&self) {
for _ in 0..200 {
match self.is_ready() {
Ok(()) => return,
Err(e) => info!("wait_ready failed, will retry (e={e})"),
}
std::thread::sleep(Duration::from_millis(100));
}
panic!("Timeout waiting Redis server ready on port {}", self.port);
}
pub fn is_ready(&self) -> Result<(), RedisError> {
let client = redis::Client::open(self.url.as_str())?;
let mut conn = client.get_connection()?;
redis::cmd("PING").query::<String>(&mut conn)?;
info!("Redis is_ready=true url={}", self.url);
Ok(())
}
pub fn url(&self) -> &str {
&self.url
}
}
impl Drop for RedisTestServer {
fn drop(&mut self) {
if let Some(mut process) = self.process.take() {
info!("Killing redis server... process={process:?}");
if let Err(e) = process.kill() {
warn!("Failed to kill redis-server process: {}", e);
}
if let Err(e) = process.wait() {
warn!("Failed to wait for redis-server process: {}", e);
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_redis_server_start_stop() {
let server = RedisTestServer::start().unwrap();
server.wait_ready();
assert!(server.url().starts_with("redis://"));
let client = redis::Client::open(server.url()).unwrap();
let mut conn = client.get_connection().unwrap();
let _: () = redis::cmd("SET")
.arg("test_key")
.arg("test_value")
.query(&mut conn)
.unwrap();
let value: String = redis::cmd("GET").arg("test_key").query(&mut conn).unwrap();
assert_eq!(value, "test_value");
}
}

134
third_party/sglang/sgl-model-gateway/tests/common/streaming_helpers.rs vendored Normal file
View File

@@ -0,0 +1,134 @@
//! Streaming Test Helpers
//!
//! Utilities for creating realistic streaming chunks that simulate
//! how LLM tokens actually arrive (1-5 characters at a time).
/// Split input into realistic char-level chunks (2-3 chars each for determinism)
pub fn create_realistic_chunks(input: &str) -> Vec<String> {
let mut chunks = Vec::new();
let chars: Vec<char> = input.chars().collect();
let mut i = 0;
while i < chars.len() {
// Take 2-3 characters at a time (deterministic for testing)
let chunk_size = if i + 3 <= chars.len() && chars[i].is_ascii_alphanumeric() {
3 // Longer chunks for alphanumeric sequences
} else {
2 // Shorter chunks for special characters
};
let end = (i + chunk_size).min(chars.len());
let chunk: String = chars[i..end].iter().collect();
chunks.push(chunk);
i = end;
}
chunks
}
/// Split input at strategic positions to test edge cases
/// This creates chunks that break at critical positions like after quotes, colons, etc.
pub fn create_strategic_chunks(input: &str) -> Vec<String> {
let mut chunks = Vec::new();
let mut current = String::new();
let chars: Vec<char> = input.chars().collect();
for (i, &ch) in chars.iter().enumerate() {
current.push(ch);
// Break after strategic characters
let should_break = matches!(ch, '"' | ':' | ',' | '{' | '}' | '[' | ']')
|| (i > 0 && chars[i-1] == '"' && ch == ' ') // Space after quote
|| current.len() >= 5; // Max 5 chars per chunk
if should_break && !current.is_empty() {
chunks.push(current.clone());
current.clear();
}
}
if !current.is_empty() {
chunks.push(current);
}
chunks
}
/// Create the bug scenario chunks: `{"name": "` arrives in parts
pub fn create_bug_scenario_chunks() -> Vec<&'static str> {
vec![
r#"{"#,
r#"""#,
r#"name"#,
r#"""#,
r#":"#,
r#" "#,
r#"""#, // Bug occurs here: parser has {"name": "
r#"search"#, // Use valid tool name
r#"""#,
r#","#,
r#" "#,
r#"""#,
r#"arguments"#,
r#"""#,
r#":"#,
r#" "#,
r#"{"#,
r#"""#,
r#"query"#,
r#"""#,
r#":"#,
r#" "#,
r#"""#,
r#"test query"#,
r#"""#,
r#"}"#,
r#"}"#,
]
}
#[cfg(test)]
mod tests {
#[allow(unused_imports)]
use super::*;
#[test]
fn test_realistic_chunks() {
let input = r#"{"name": "test"}"#;
let chunks = create_realistic_chunks(input);
// Should have multiple chunks
assert!(chunks.len() > 3);
// Reconstructed should equal original
let reconstructed: String = chunks.join("");
assert_eq!(reconstructed, input);
}
#[test]
fn test_strategic_chunks_breaks_after_quotes() {
let input = r#"{"name": "value"}"#;
let chunks = create_strategic_chunks(input);
// Should break after quotes and colons
assert!(chunks.iter().any(|c| c.ends_with('"')));
assert!(chunks.iter().any(|c| c.ends_with(':')));
// Reconstructed should equal original
let reconstructed: String = chunks.join("");
assert_eq!(reconstructed, input);
}
#[test]
fn test_bug_scenario_chunks() {
let chunks = create_bug_scenario_chunks();
let reconstructed: String = chunks.join("");
// Should reconstruct to valid JSON
assert!(reconstructed.contains(r#"{"name": "search""#));
// The critical chunk sequence should be present (space after colon, then quote in next chunk)
let joined = chunks.join("|");
assert!(joined.contains(r#" |"#)); // The bug happens at {"name": " and then "
}
}

267
third_party/sglang/sgl-model-gateway/tests/common/test_app.rs vendored Normal file
View File

@@ -0,0 +1,267 @@
use std::sync::{Arc, OnceLock};
use axum::Router;
use data_connector::{
MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage,
};
use reqwest::Client;
use smg::{
app_context::AppContext,
config::RouterConfig,
core::{
BasicWorkerBuilder, LoadMonitor, ModelCard, RuntimeType, Worker, WorkerRegistry, WorkerType,
},
middleware::{AuthConfig, TokenBucket},
policies::PolicyRegistry,
routers::RouterTrait,
server::{build_app, AppState},
tokenizer::registry::TokenizerRegistry,
};
use smg_mcp::{McpConfig, McpManager};
/// Create a test Axum application using the actual server's build_app function
#[allow(dead_code)]
pub fn create_test_app(
router: Arc<dyn RouterTrait>,
client: Client,
router_config: &RouterConfig,
) -> Router {
// Initialize rate limiter
let rate_limiter = match router_config.max_concurrent_requests {
n if n <= 0 => None,
n => {
let rate_limit_tokens = router_config
.rate_limit_tokens_per_second
.filter(|&t| t > 0)
.unwrap_or(n);
Some(Arc::new(TokenBucket::new(
n as usize,
rate_limit_tokens as usize,
)))
}
};
// Initialize registries
let worker_registry = Arc::new(WorkerRegistry::new());
let policy_registry = Arc::new(PolicyRegistry::new(router_config.policy.clone()));
// Initialize storage backends
let response_storage = Arc::new(MemoryResponseStorage::new());
let conversation_storage = Arc::new(MemoryConversationStorage::new());
let conversation_item_storage = Arc::new(MemoryConversationItemStorage::new());
// Initialize load monitor
let load_monitor = Some(Arc::new(LoadMonitor::new(
worker_registry.clone(),
policy_registry.clone(),
client.clone(),
router_config.worker_startup_check_interval_secs,
)));
// Create empty OnceLock for worker job queue and workflow engines
let worker_job_queue = Arc::new(OnceLock::new());
let workflow_engines = Arc::new(OnceLock::new());
// Create AppContext using builder pattern
let app_context = Arc::new(
AppContext::builder()
.router_config(router_config.clone())
.client(client)
.rate_limiter(rate_limiter)
.tokenizer_registry(Arc::new(TokenizerRegistry::new())) // tokenizer
.reasoning_parser_factory(None) // reasoning_parser_factory
.tool_parser_factory(None) // tool_parser_factory
.worker_registry(worker_registry)
.policy_registry(policy_registry)
.response_storage(response_storage)
.conversation_storage(conversation_storage)
.conversation_item_storage(conversation_item_storage)
.load_monitor(load_monitor)
.worker_job_queue(worker_job_queue)
.workflow_engines(workflow_engines)
.build()
.unwrap(),
);
// Create AppState with the test router and context
let app_state = Arc::new(AppState {
router,
context: app_context,
concurrency_queue_tx: None,
router_manager: None,
mesh_handler: None,
mesh_sync_manager: None,
});
// Configure request ID headers (use defaults if not specified)
let request_id_headers = router_config.request_id_headers.clone().unwrap_or_else(|| {
vec![
"x-request-id".to_string(),
"x-correlation-id".to_string(),
"x-trace-id".to_string(),
"request-id".to_string(),
]
});
// Create auth config from router config
let auth_config = AuthConfig {
api_key: router_config.api_key.clone(),
};
// Use the actual server's build_app function
build_app(
app_state,
auth_config,
None, // No control plane auth for tests
router_config.max_payload_size,
request_id_headers,
router_config.cors_allowed_origins.clone(),
)
}
/// Create a test Axum application with an existing AppContext
#[allow(dead_code)]
pub fn create_test_app_with_context(
router: Arc<dyn RouterTrait>,
app_context: Arc<AppContext>,
) -> Router {
// Create AppState with the test router and context
let app_state = Arc::new(AppState {
router,
context: app_context.clone(),
concurrency_queue_tx: None,
router_manager: None,
mesh_handler: None,
mesh_sync_manager: None,
});
// Get config from the context
let router_config = &app_context.router_config;
// Configure request ID headers (use defaults if not specified)
let request_id_headers = router_config.request_id_headers.clone().unwrap_or_else(|| {
vec![
"x-request-id".to_string(),
"x-correlation-id".to_string(),
"x-trace-id".to_string(),
"request-id".to_string(),
]
});
// Create auth config from router config
let auth_config = AuthConfig {
api_key: router_config.api_key.clone(),
};
// Use the actual server's build_app function
build_app(
app_state,
auth_config,
None, // No control plane auth for tests
router_config.max_payload_size,
request_id_headers,
router_config.cors_allowed_origins.clone(),
)
}
/// Create a minimal test AppContext for unit tests
#[allow(dead_code)]
pub async fn create_test_app_context() -> Arc<AppContext> {
let router_config = RouterConfig::default();
let client = Client::new();
// Initialize empty OnceLocks
let worker_job_queue = Arc::new(OnceLock::new());
let workflow_engines = Arc::new(OnceLock::new());
// Initialize MCP manager with empty config
let mcp_manager_lock = Arc::new(OnceLock::new());
let empty_config = McpConfig {
servers: vec![],
pool: Default::default(),
proxy: None,
warmup: vec![],
inventory: Default::default(),
};
let mcp_manager = McpManager::with_defaults(empty_config)
.await
.expect("Failed to create MCP manager");
mcp_manager_lock.set(Arc::new(mcp_manager)).ok();
// Initialize registries
let worker_registry = Arc::new(WorkerRegistry::new());
let policy_registry = Arc::new(PolicyRegistry::new(router_config.policy.clone()));
// Initialize storage backends
let response_storage = Arc::new(MemoryResponseStorage::new());
let conversation_storage = Arc::new(MemoryConversationStorage::new());
let conversation_item_storage = Arc::new(MemoryConversationItemStorage::new());
Arc::new(
AppContext::builder()
.router_config(router_config)
.client(client)
.rate_limiter(None)
.tokenizer_registry(Arc::new(TokenizerRegistry::new()))
.reasoning_parser_factory(None)
.tool_parser_factory(None)
.worker_registry(worker_registry)
.policy_registry(policy_registry)
.response_storage(response_storage)
.conversation_storage(conversation_storage)
.conversation_item_storage(conversation_item_storage)
.load_monitor(None)
.worker_job_queue(worker_job_queue)
.workflow_engines(workflow_engines)
.mcp_manager(mcp_manager_lock)
.build()
.unwrap(),
)
}
/// Register an external worker (OpenAI-compatible API endpoint) in the test AppContext.
///
/// This is used by tests that need to test the OpenAI router, which expects
/// workers to be registered in the WorkerRegistry before routing requests.
///
/// # Arguments
/// * `ctx` - The AppContext to register the worker in
/// * `url` - The base URL of the external API endpoint
/// * `models` - Optional list of model IDs this worker supports. If empty, uses "gpt-3.5-turbo" as default.
#[allow(dead_code)]
pub fn register_external_worker(ctx: &Arc<AppContext>, url: &str, models: Option<Vec<&str>>) {
let model_list: Vec<ModelCard> = models
.unwrap_or_else(|| vec!["gpt-3.5-turbo"])
.into_iter()
.map(ModelCard::new)
.collect();
let worker: Arc<dyn Worker> = Arc::new(
BasicWorkerBuilder::new(url)
.worker_type(WorkerType::Regular)
.runtime_type(RuntimeType::External)
.models(model_list)
.build(),
);
ctx.worker_registry.register(worker);
}
/// Register an external worker with a custom model card that has aliases.
///
/// # Arguments
/// * `ctx` - The AppContext to register the worker in
/// * `url` - The base URL of the external API endpoint
/// * `model_card` - A fully configured ModelCard with aliases, provider, etc.
#[allow(dead_code)]
pub fn register_external_worker_with_card(ctx: &Arc<AppContext>, url: &str, model_card: ModelCard) {
let worker: Arc<dyn Worker> = Arc::new(
BasicWorkerBuilder::new(url)
.worker_type(WorkerType::Regular)
.runtime_type(RuntimeType::External)
.model(model_card)
.build(),
);
ctx.worker_registry.register(worker);
}

299
third_party/sglang/sgl-model-gateway/tests/common/test_certs.rs vendored Normal file
View File

@@ -0,0 +1,299 @@
// TLS certificate generation for integration tests
#![allow(dead_code)]
use std::path::PathBuf;
use openssl::{
asn1::Asn1Time,
bn::{BigNum, MsbOption},
hash::MessageDigest,
pkey::{PKey, Private},
rsa::Rsa,
x509::{
extension::{BasicConstraints, ExtendedKeyUsage, KeyUsage, SubjectAlternativeName},
X509Builder, X509NameBuilder, X509,
},
};
use tempfile::TempDir;
/// Container for generated test certificates
pub struct TestCertificates {
/// Temporary directory containing the certificate files
pub temp_dir: TempDir,
/// Path to CA certificate
pub ca_cert_path: PathBuf,
/// Path to CA private key
pub ca_key_path: PathBuf,
/// Path to server certificate
pub server_cert_path: PathBuf,
/// Path to server private key
pub server_key_path: PathBuf,
/// Path to client certificate
pub client_cert_path: PathBuf,
/// Path to client private key
pub client_key_path: PathBuf,
}
impl TestCertificates {
/// Generate a complete set of test certificates for mTLS testing.
///
/// Creates:
/// - CA certificate and key (self-signed root CA)
/// - Server certificate and key (signed by CA, for localhost)
/// - Client certificate and key (signed by CA, for client authentication)
pub fn generate() -> Result<Self, Box<dyn std::error::Error>> {
let temp_dir = TempDir::new()?;
let base_path = temp_dir.path();
// Generate CA key pair
let ca_key = generate_rsa_key()?;
let ca_cert = generate_ca_certificate(&ca_key)?;
// Generate server key pair and certificate
let server_key = generate_rsa_key()?;
let server_cert = generate_server_certificate(&server_key, &ca_cert, &ca_key)?;
// Generate client key pair and certificate
let client_key = generate_rsa_key()?;
let client_cert = generate_client_certificate(&client_key, &ca_cert, &ca_key)?;
// Write all files
let ca_cert_path = base_path.join("ca_cert.pem");
let ca_key_path = base_path.join("ca_key.pem");
let server_cert_path = base_path.join("server_cert.pem");
let server_key_path = base_path.join("server_key.pem");
let client_cert_path = base_path.join("client_cert.pem");
let client_key_path = base_path.join("client_key.pem");
std::fs::write(&ca_cert_path, ca_cert.to_pem()?)?;
std::fs::write(&ca_key_path, ca_key.private_key_to_pem_pkcs8()?)?;
std::fs::write(&server_cert_path, server_cert.to_pem()?)?;
std::fs::write(&server_key_path, server_key.private_key_to_pem_pkcs8()?)?;
std::fs::write(&client_cert_path, client_cert.to_pem()?)?;
std::fs::write(&client_key_path, client_key.private_key_to_pem_pkcs8()?)?;
Ok(Self {
temp_dir,
ca_cert_path,
ca_key_path,
server_cert_path,
server_key_path,
client_cert_path,
client_key_path,
})
}
/// Get paths as string references for use with RouterConfig builder
pub fn ca_cert_str(&self) -> &str {
self.ca_cert_path.to_str().unwrap()
}
pub fn server_cert_str(&self) -> &str {
self.server_cert_path.to_str().unwrap()
}
pub fn server_key_str(&self) -> &str {
self.server_key_path.to_str().unwrap()
}
pub fn client_cert_str(&self) -> &str {
self.client_cert_path.to_str().unwrap()
}
pub fn client_key_str(&self) -> &str {
self.client_key_path.to_str().unwrap()
}
}
/// Generate a 2048-bit RSA key pair
fn generate_rsa_key() -> Result<PKey<Private>, Box<dyn std::error::Error>> {
let rsa = Rsa::generate(2048)?;
Ok(PKey::from_rsa(rsa)?)
}
/// Generate a self-signed CA certificate
fn generate_ca_certificate(key: &PKey<Private>) -> Result<X509, Box<dyn std::error::Error>> {
let mut name_builder = X509NameBuilder::new()?;
name_builder.append_entry_by_text("C", "US")?;
name_builder.append_entry_by_text("ST", "California")?;
name_builder.append_entry_by_text("L", "Test City")?;
name_builder.append_entry_by_text("O", "Test CA Organization")?;
name_builder.append_entry_by_text("CN", "Test CA")?;
let name = name_builder.build();
let mut cert_builder = X509Builder::new()?;
cert_builder.set_version(2)?; // X509 v3
// Serial number
let serial = {
let mut bn = BigNum::new()?;
bn.rand(128, MsbOption::MAYBE_ZERO, false)?;
bn.to_asn1_integer()?
};
cert_builder.set_serial_number(&serial)?;
cert_builder.set_subject_name(&name)?;
cert_builder.set_issuer_name(&name)?; // Self-signed
cert_builder.set_pubkey(key)?;
// Validity: 1 year from now
let not_before = Asn1Time::days_from_now(0)?;
let not_after = Asn1Time::days_from_now(365)?;
cert_builder.set_not_before(&not_before)?;
cert_builder.set_not_after(&not_after)?;
// Extensions for CA
let basic_constraints = BasicConstraints::new().critical().ca().build()?;
cert_builder.append_extension(basic_constraints)?;
let key_usage = KeyUsage::new()
.critical()
.key_cert_sign()
.crl_sign()
.build()?;
cert_builder.append_extension(key_usage)?;
cert_builder.sign(key, MessageDigest::sha256())?;
Ok(cert_builder.build())
}
/// Generate a server certificate signed by the CA
fn generate_server_certificate(
key: &PKey<Private>,
ca_cert: &X509,
ca_key: &PKey<Private>,
) -> Result<X509, Box<dyn std::error::Error>> {
let mut name_builder = X509NameBuilder::new()?;
name_builder.append_entry_by_text("C", "US")?;
name_builder.append_entry_by_text("ST", "California")?;
name_builder.append_entry_by_text("L", "Test City")?;
name_builder.append_entry_by_text("O", "Test Server Organization")?;
name_builder.append_entry_by_text("CN", "localhost")?;
let name = name_builder.build();
let mut cert_builder = X509Builder::new()?;
cert_builder.set_version(2)?;
let serial = {
let mut bn = BigNum::new()?;
bn.rand(128, MsbOption::MAYBE_ZERO, false)?;
bn.to_asn1_integer()?
};
cert_builder.set_serial_number(&serial)?;
cert_builder.set_subject_name(&name)?;
cert_builder.set_issuer_name(ca_cert.subject_name())?;
cert_builder.set_pubkey(key)?;
let not_before = Asn1Time::days_from_now(0)?;
let not_after = Asn1Time::days_from_now(365)?;
cert_builder.set_not_before(&not_before)?;
cert_builder.set_not_after(&not_after)?;
// Extensions for server certificate
let basic_constraints = BasicConstraints::new().build()?;
cert_builder.append_extension(basic_constraints)?;
let key_usage = KeyUsage::new()
.critical()
.digital_signature()
.key_encipherment()
.build()?;
cert_builder.append_extension(key_usage)?;
let ext_key_usage = ExtendedKeyUsage::new().server_auth().build()?;
cert_builder.append_extension(ext_key_usage)?;
// Subject Alternative Names for localhost
let san = SubjectAlternativeName::new()
.dns("localhost")
.ip("127.0.0.1")
.ip("::1")
.build(&cert_builder.x509v3_context(Some(ca_cert), None))?;
cert_builder.append_extension(san)?;
cert_builder.sign(ca_key, MessageDigest::sha256())?;
Ok(cert_builder.build())
}
/// Generate a client certificate signed by the CA
fn generate_client_certificate(
key: &PKey<Private>,
ca_cert: &X509,
ca_key: &PKey<Private>,
) -> Result<X509, Box<dyn std::error::Error>> {
let mut name_builder = X509NameBuilder::new()?;
name_builder.append_entry_by_text("C", "US")?;
name_builder.append_entry_by_text("ST", "California")?;
name_builder.append_entry_by_text("L", "Test City")?;
name_builder.append_entry_by_text("O", "Test Client Organization")?;
name_builder.append_entry_by_text("CN", "Test Client")?;
let name = name_builder.build();
let mut cert_builder = X509Builder::new()?;
cert_builder.set_version(2)?;
let serial = {
let mut bn = BigNum::new()?;
bn.rand(128, MsbOption::MAYBE_ZERO, false)?;
bn.to_asn1_integer()?
};
cert_builder.set_serial_number(&serial)?;
cert_builder.set_subject_name(&name)?;
cert_builder.set_issuer_name(ca_cert.subject_name())?;
cert_builder.set_pubkey(key)?;
let not_before = Asn1Time::days_from_now(0)?;
let not_after = Asn1Time::days_from_now(365)?;
cert_builder.set_not_before(&not_before)?;
cert_builder.set_not_after(&not_after)?;
// Extensions for client certificate
let basic_constraints = BasicConstraints::new().build()?;
cert_builder.append_extension(basic_constraints)?;
let key_usage = KeyUsage::new().critical().digital_signature().build()?;
cert_builder.append_extension(key_usage)?;
let ext_key_usage = ExtendedKeyUsage::new().client_auth().build()?;
cert_builder.append_extension(ext_key_usage)?;
cert_builder.sign(ca_key, MessageDigest::sha256())?;
Ok(cert_builder.build())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_certificate_generation() {
let certs = TestCertificates::generate().expect("Failed to generate certificates");
// Verify all files exist
assert!(certs.ca_cert_path.exists(), "CA cert should exist");
assert!(certs.ca_key_path.exists(), "CA key should exist");
assert!(certs.server_cert_path.exists(), "Server cert should exist");
assert!(certs.server_key_path.exists(), "Server key should exist");
assert!(certs.client_cert_path.exists(), "Client cert should exist");
assert!(certs.client_key_path.exists(), "Client key should exist");
// Verify files are not empty
assert!(
std::fs::metadata(&certs.ca_cert_path).unwrap().len() > 0,
"CA cert should not be empty"
);
assert!(
std::fs::metadata(&certs.server_cert_path).unwrap().len() > 0,
"Server cert should not be empty"
);
assert!(
std::fs::metadata(&certs.client_cert_path).unwrap().len() > 0,
"Client cert should not be empty"
);
}
}

336
third_party/sglang/sgl-model-gateway/tests/common/test_config.rs vendored Normal file
View File

@@ -0,0 +1,336 @@
//! Test configuration builders to reduce duplication across tests
//!
//! Provides pre-configured RouterConfig and MockWorkerConfig builders
//! for common test scenarios.
use smg::config::{
CircuitBreakerConfig, ManualAssignmentMode, PolicyConfig, RetryConfig, RouterConfig,
};
use super::mock_worker::{HealthStatus, MockWorkerConfig, WorkerType};
/// Default test configuration values
pub mod defaults {
pub const HOST: &str = "127.0.0.1";
pub const MAX_PAYLOAD_SIZE: usize = 256 * 1024 * 1024; // 256MB
pub const REQUEST_TIMEOUT_SECS: u64 = 600;
pub const WORKER_STARTUP_TIMEOUT_SECS: u64 = 5;
pub const WORKER_STARTUP_CHECK_INTERVAL_SECS: u64 = 1;
pub const MAX_CONCURRENT_REQUESTS: i32 = 64;
pub const QUEUE_TIMEOUT_SECS: u64 = 60;
}
/// Builder for common test RouterConfig patterns
pub struct TestRouterConfig;
impl TestRouterConfig {
/// Create a basic round-robin config for routing tests
pub fn round_robin(port: u16) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.build_unchecked()
}
/// Create a random load balancing config
pub fn random(port: u16) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.random_policy()
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.build_unchecked()
}
/// Create a cache-aware config for routing tests
pub fn cache_aware(port: u16) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.cache_aware_policy(
0.5, // cache_threshold
32, // balance_abs_threshold
1.5, // balance_rel_threshold
60, // eviction_interval_secs
1000, // max_tree_size
)
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.build_unchecked()
}
/// Create a power-of-two config
pub fn power_of_two(port: u16) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.power_of_two_policy(5) // load_check_interval_secs
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.build_unchecked()
}
/// Create a manual routing config (for sticky routing tests)
pub fn manual(port: u16) -> RouterConfig {
Self::manual_with_mode(port, ManualAssignmentMode::Random)
}
/// Create a manual routing config with min_group assignment mode
pub fn manual_min_group(port: u16) -> RouterConfig {
Self::manual_with_mode(port, ManualAssignmentMode::MinGroup)
}
/// Create a manual routing config with specified assignment mode
pub fn manual_with_mode(port: u16, assignment_mode: ManualAssignmentMode) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.policy(PolicyConfig::Manual {
eviction_interval_secs: 60,
max_idle_secs: 3600,
assignment_mode,
})
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.build_unchecked()
}
/// Create a config with custom concurrent request limit (for rate limiting tests)
pub fn with_concurrency(port: u16, max_concurrent: i32) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(max_concurrent)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.build_unchecked()
}
/// Create a config with custom payload size limit
pub fn with_payload_limit(port: u16, max_payload_size: usize) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host(defaults::HOST)
.port(port)
.max_payload_size(max_payload_size)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.build_unchecked()
}
/// Create a config with short timeouts (for timeout/retry tests)
pub fn with_short_timeouts(port: u16) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(5)
.worker_startup_timeout_secs(2)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(5)
.build_unchecked()
}
/// Create a round-robin config with retry settings
pub fn round_robin_with_retry(port: u16, retry_config: RetryConfig) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.retry_config(retry_config)
.build_unchecked()
}
/// Create a round-robin config with circuit breaker
pub fn round_robin_with_circuit_breaker(
port: u16,
circuit_breaker: CircuitBreakerConfig,
) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.circuit_breaker_config(circuit_breaker)
.build_unchecked()
}
/// Create a round-robin config with both retry and circuit breaker
pub fn round_robin_with_reliability(
port: u16,
retry_config: RetryConfig,
circuit_breaker: CircuitBreakerConfig,
) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.retry_config(retry_config)
.circuit_breaker_config(circuit_breaker)
.build_unchecked()
}
}
/// Builder for common MockWorkerConfig patterns
pub struct TestWorkerConfig;
impl TestWorkerConfig {
/// Create a healthy worker config
pub fn healthy(port: u16) -> MockWorkerConfig {
MockWorkerConfig {
port,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}
}
/// Create multiple healthy workers with sequential ports
pub fn healthy_workers(start_port: u16, count: u16) -> Vec<MockWorkerConfig> {
(0..count).map(|i| Self::healthy(start_port + i)).collect()
}
/// Create an unhealthy worker config
pub fn unhealthy(port: u16) -> MockWorkerConfig {
MockWorkerConfig {
port,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Unhealthy,
response_delay_ms: 0,
fail_rate: 0.0,
}
}
/// Create a slow worker config (for timeout tests)
pub fn slow(port: u16, delay_ms: u64) -> MockWorkerConfig {
MockWorkerConfig {
port,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: delay_ms,
fail_rate: 0.0,
}
}
/// Create multiple slow workers with sequential ports
pub fn slow_workers(start_port: u16, count: u16, delay_ms: u64) -> Vec<MockWorkerConfig> {
(0..count)
.map(|i| Self::slow(start_port + i, delay_ms))
.collect()
}
/// Create a flaky worker config (for retry/fault tolerance tests)
pub fn flaky(port: u16, fail_rate: f32) -> MockWorkerConfig {
MockWorkerConfig {
port,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate,
}
}
/// Create a decode worker config (for PD routing tests)
pub fn decode(port: u16) -> MockWorkerConfig {
MockWorkerConfig {
port,
worker_type: WorkerType::Decode,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}
}
/// Create a prefill worker config (for PD routing tests)
pub fn prefill(port: u16) -> MockWorkerConfig {
MockWorkerConfig {
port,
worker_type: WorkerType::Prefill,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_round_robin_config() {
let config = TestRouterConfig::round_robin(3000);
assert_eq!(config.port, 3000);
}
#[test]
fn test_healthy_workers() {
let workers = TestWorkerConfig::healthy_workers(8000, 3);
assert_eq!(workers.len(), 3);
assert_eq!(workers[0].port, 8000);
assert_eq!(workers[1].port, 8001);
assert_eq!(workers[2].port, 8002);
}
}

370
third_party/sglang/sgl-model-gateway/tests/common/tls_mock_worker.rs vendored Normal file
View File

@@ -0,0 +1,370 @@
// TLS-enabled mock worker for mTLS integration tests
#![allow(dead_code)]
use std::{
net::SocketAddr,
path::Path,
sync::{Arc, Once},
time::{SystemTime, UNIX_EPOCH},
};
use axum::{
extract::{Json, State},
http::StatusCode,
response::{IntoResponse, Response},
routing::{get, post},
Router,
};
use axum_server::tls_rustls::RustlsConfig;
use rustls::{server::WebPkiClientVerifier, RootCertStore};
use serde_json::json;
use tokio::sync::RwLock;
// Ensure crypto provider is installed exactly once
static CRYPTO_PROVIDER_INIT: Once = Once::new();
fn ensure_crypto_provider() {
CRYPTO_PROVIDER_INIT.call_once(|| {
let _ = rustls::crypto::ring::default_provider().install_default();
});
}
/// Configuration for TLS mock worker behavior
#[derive(Clone)]
pub struct TlsMockWorkerConfig {
pub port: u16,
/// Require client certificate (mTLS) or just server TLS
pub require_client_cert: bool,
/// Response delay in milliseconds
pub response_delay_ms: u64,
/// Fail rate (0.0 - 1.0)
pub fail_rate: f32,
}
impl Default for TlsMockWorkerConfig {
fn default() -> Self {
Self {
port: 0,
require_client_cert: true,
response_delay_ms: 0,
fail_rate: 0.0,
}
}
}
/// TLS-enabled mock worker server for mTLS testing
pub struct TlsMockWorker {
config: Arc<RwLock<TlsMockWorkerConfig>>,
shutdown_handle: Option<tokio::task::JoinHandle<()>>,
shutdown_tx: Option<tokio::sync::oneshot::Sender<()>>,
}
impl TlsMockWorker {
pub fn new(config: TlsMockWorkerConfig) -> Self {
Self {
config: Arc::new(RwLock::new(config)),
shutdown_handle: None,
shutdown_tx: None,
}
}
/// Start the TLS mock worker server
///
/// # Arguments
/// * `server_cert_path` - Path to server certificate PEM file
/// * `server_key_path` - Path to server private key PEM file
/// * `ca_cert_path` - Path to CA certificate for client verification (optional for TLS-only mode)
pub async fn start(
&mut self,
server_cert_path: &Path,
server_key_path: &Path,
ca_cert_path: Option<&Path>,
) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
// Ensure crypto provider is installed before using rustls
ensure_crypto_provider();
let config = self.config.clone();
let port = config.read().await.port;
let require_client_cert = config.read().await.require_client_cert;
// If port is 0, find an available port
let port = if port == 0 {
let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
let port = listener.local_addr()?.port();
drop(listener);
config.write().await.port = port;
port
} else {
port
};
let app = Router::new()
.route("/health", get(health_handler))
.route("/health_generate", get(health_generate_handler))
.route("/server_info", get(server_info_handler))
.route("/generate", post(generate_handler))
.route("/v1/chat/completions", post(chat_completions_handler))
.with_state(config);
let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel::<()>();
self.shutdown_tx = Some(shutdown_tx);
// Build TLS configuration
let rustls_config = if require_client_cert {
// mTLS: require client certificate
let ca_cert_path = ca_cert_path.ok_or("CA cert path required for mTLS")?;
build_mtls_config(server_cert_path, server_key_path, ca_cert_path).await?
} else {
// TLS only: no client cert required
build_tls_config(server_cert_path, server_key_path).await?
};
let addr = SocketAddr::from(([127, 0, 0, 1], port));
// Spawn the server in a separate task
let handle = tokio::spawn(async move {
let server =
axum_server::bind_rustls(addr, rustls_config).serve(app.into_make_service());
tokio::select! {
result = server => {
if let Err(e) = result {
eprintln!("TLS Server error: {}", e);
}
}
_ = &mut shutdown_rx => {
// Graceful shutdown
}
}
});
self.shutdown_handle = Some(handle);
// Wait for the server to start
tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;
let url = format!("https://127.0.0.1:{}", port);
Ok(url)
}
/// Stop the TLS mock worker server
pub async fn stop(&mut self) {
if let Some(shutdown_tx) = self.shutdown_tx.take() {
let _ = shutdown_tx.send(());
}
if let Some(handle) = self.shutdown_handle.take() {
let _ = tokio::time::timeout(tokio::time::Duration::from_secs(5), handle).await;
}
}
}
impl Drop for TlsMockWorker {
fn drop(&mut self) {
if let Some(shutdown_tx) = self.shutdown_tx.take() {
let _ = shutdown_tx.send(());
}
}
}
/// Build TLS config for server-only TLS (no client cert required)
async fn build_tls_config(
cert_path: &Path,
key_path: &Path,
) -> Result<RustlsConfig, Box<dyn std::error::Error + Send + Sync>> {
let config = RustlsConfig::from_pem_file(cert_path, key_path).await?;
Ok(config)
}
/// Build mTLS config requiring client certificate
async fn build_mtls_config(
cert_path: &Path,
key_path: &Path,
ca_cert_path: &Path,
) -> Result<RustlsConfig, Box<dyn std::error::Error + Send + Sync>> {
use std::io::BufReader;
use rustls_pemfile::{certs, pkcs8_private_keys};
// Read server certificate
let cert_file = std::fs::File::open(cert_path)?;
let mut reader = BufReader::new(cert_file);
let cert_chain: Vec<_> = certs(&mut reader).filter_map(|r| r.ok()).collect();
// Read server private key
let key_file = std::fs::File::open(key_path)?;
let mut reader = BufReader::new(key_file);
let private_key = pkcs8_private_keys(&mut reader)
.next()
.ok_or("No private key found")??;
// Read CA certificate for client verification
let ca_file = std::fs::File::open(ca_cert_path)?;
let mut reader = BufReader::new(ca_file);
let ca_certs: Vec<_> = certs(&mut reader).filter_map(|r| r.ok()).collect();
// Build root certificate store for client verification
let mut root_store = RootCertStore::empty();
for cert in ca_certs {
root_store.add(cert)?;
}
// Create client certificate verifier
let client_verifier = WebPkiClientVerifier::builder(Arc::new(root_store))
.build()
.map_err(|e| format!("Failed to build client verifier: {}", e))?;
// Build server config with client verification
let server_config = rustls::ServerConfig::builder()
.with_client_cert_verifier(client_verifier)
.with_single_cert(cert_chain, private_key.into())
.map_err(|e| format!("Failed to build server config: {}", e))?;
Ok(RustlsConfig::from_config(Arc::new(server_config)))
}
// Handler implementations (simplified versions of mock_worker handlers)
async fn should_fail(config: &TlsMockWorkerConfig) -> bool {
rand::random::<f32>() < config.fail_rate
}
async fn health_handler(State(config): State<Arc<RwLock<TlsMockWorkerConfig>>>) -> Response {
let config = config.read().await;
if should_fail(&config).await {
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({ "error": "Random failure" })),
)
.into_response();
}
Json(json!({
"status": "healthy",
"timestamp": SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(),
"tls_enabled": true
}))
.into_response()
}
async fn health_generate_handler(
State(config): State<Arc<RwLock<TlsMockWorkerConfig>>>,
) -> Response {
let config = config.read().await;
if should_fail(&config).await {
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({ "error": "Random failure" })),
)
.into_response();
}
Json(json!({
"status": "ok",
"queue_length": 0,
"processing_time_ms": config.response_delay_ms
}))
.into_response()
}
async fn server_info_handler(State(config): State<Arc<RwLock<TlsMockWorkerConfig>>>) -> Response {
let config = config.read().await;
if should_fail(&config).await {
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({ "error": "Random failure" })),
)
.into_response();
}
Json(json!({
"model_path": "mock-tls-model",
"port": config.port,
"host": "127.0.0.1",
"tls_enabled": true,
"version": "0.3.0"
}))
.into_response()
}
async fn generate_handler(
State(config): State<Arc<RwLock<TlsMockWorkerConfig>>>,
Json(_payload): Json<serde_json::Value>,
) -> Response {
let config = config.read().await;
if should_fail(&config).await {
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({ "error": "Random failure" })),
)
.into_response();
}
if config.response_delay_ms > 0 {
tokio::time::sleep(tokio::time::Duration::from_millis(config.response_delay_ms)).await;
}
Json(json!({
"text": "This is a mock TLS response.",
"meta_info": {
"prompt_tokens": 10,
"completion_tokens": 5,
"tls_verified": true
}
}))
.into_response()
}
async fn chat_completions_handler(
State(config): State<Arc<RwLock<TlsMockWorkerConfig>>>,
Json(_payload): Json<serde_json::Value>,
) -> Response {
let config = config.read().await;
if should_fail(&config).await {
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({
"error": {
"message": "Random failure",
"type": "internal_error"
}
})),
)
.into_response();
}
if config.response_delay_ms > 0 {
tokio::time::sleep(tokio::time::Duration::from_millis(config.response_delay_ms)).await;
}
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
Json(json!({
"id": format!("chatcmpl-{}", uuid::Uuid::new_v4()),
"object": "chat.completion",
"created": timestamp,
"model": "mock-tls-model",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": "This is a mock TLS chat response."
},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 5,
"total_tokens": 15
}
}))
.into_response()
}