1310 lines
42 KiB
Rust
1310 lines
42 KiB
Rust
mod common;
|
|
|
|
use actix_web::{http::StatusCode, rt::System, test as actix_test, web, App};
|
|
use common::mock_worker::{HealthStatus, MockWorker, MockWorkerConfig, WorkerType};
|
|
use reqwest::Client;
|
|
use serde_json::json;
|
|
use sglang_router_rs::config::{PolicyConfig, RouterConfig, RoutingMode};
|
|
use sglang_router_rs::server::{
|
|
add_worker, flush_cache, generate, get_loads, get_model_info, get_server_info, health,
|
|
health_generate, list_workers, liveness, readiness, remove_worker, v1_chat_completions,
|
|
v1_completions, v1_models, AppState,
|
|
};
|
|
|
|
/// Test context that manages mock workers
|
|
struct TestContext {
|
|
workers: Vec<MockWorker>,
|
|
app_state: web::Data<AppState>,
|
|
}
|
|
|
|
impl TestContext {
|
|
async fn new(worker_configs: Vec<MockWorkerConfig>) -> Self {
|
|
// Create default router config
|
|
let config = RouterConfig {
|
|
mode: RoutingMode::Regular {
|
|
worker_urls: vec![],
|
|
},
|
|
policy: PolicyConfig::Random,
|
|
host: "127.0.0.1".to_string(),
|
|
port: 3002,
|
|
max_payload_size: 256 * 1024 * 1024,
|
|
request_timeout_secs: 600,
|
|
worker_startup_timeout_secs: 1,
|
|
worker_startup_check_interval_secs: 1,
|
|
discovery: None,
|
|
metrics: None,
|
|
log_dir: None,
|
|
log_level: None,
|
|
};
|
|
|
|
Self::new_with_config(config, worker_configs).await
|
|
}
|
|
|
|
async fn new_with_config(config: RouterConfig, worker_configs: Vec<MockWorkerConfig>) -> Self {
|
|
let mut workers = Vec::new();
|
|
let mut worker_urls = Vec::new();
|
|
|
|
// Start mock workers if any
|
|
for worker_config in worker_configs {
|
|
let mut worker = MockWorker::new(worker_config);
|
|
let url = worker.start().await.unwrap();
|
|
worker_urls.push(url);
|
|
workers.push(worker);
|
|
}
|
|
|
|
if !workers.is_empty() {
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;
|
|
}
|
|
|
|
let client = Client::builder()
|
|
.timeout(std::time::Duration::from_secs(config.request_timeout_secs))
|
|
.build()
|
|
.unwrap();
|
|
|
|
let app_state = AppState::new(config, client).unwrap();
|
|
let app_state = web::Data::new(app_state);
|
|
|
|
// Add workers if any
|
|
if !worker_urls.is_empty() {
|
|
let app = actix_test::init_service(
|
|
App::new().app_data(app_state.clone()).service(add_worker),
|
|
)
|
|
.await;
|
|
|
|
for url in &worker_urls {
|
|
let req = actix_test::TestRequest::post()
|
|
.uri(&format!("/add_worker?url={}", url))
|
|
.to_request();
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert!(resp.status().is_success());
|
|
}
|
|
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
|
|
}
|
|
|
|
Self { workers, app_state }
|
|
}
|
|
|
|
async fn create_app(
|
|
&self,
|
|
) -> impl actix_web::dev::Service<
|
|
actix_http::Request,
|
|
Response = actix_web::dev::ServiceResponse,
|
|
Error = actix_web::Error,
|
|
> {
|
|
actix_test::init_service(
|
|
App::new()
|
|
.app_data(self.app_state.clone())
|
|
.service(liveness)
|
|
.service(readiness)
|
|
.service(health)
|
|
.service(health_generate)
|
|
.service(get_server_info)
|
|
.service(get_model_info)
|
|
.service(v1_models)
|
|
.service(generate)
|
|
.service(v1_chat_completions)
|
|
.service(v1_completions)
|
|
.service(add_worker)
|
|
.service(list_workers)
|
|
.service(remove_worker)
|
|
.service(flush_cache)
|
|
.service(get_loads),
|
|
)
|
|
.await
|
|
}
|
|
|
|
async fn shutdown(mut self) {
|
|
for worker in &mut self.workers {
|
|
worker.stop().await;
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod health_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_liveness_endpoint() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![]).await;
|
|
let app = ctx.create_app().await;
|
|
|
|
let req = actix_test::TestRequest::get().uri("/liveness").to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_readiness_with_healthy_workers() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![MockWorkerConfig {
|
|
port: 18001,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let req = actix_test::TestRequest::get()
|
|
.uri("/readiness")
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_readiness_with_unhealthy_workers() {
|
|
System::new().block_on(async {
|
|
// Create an empty context (no workers)
|
|
let ctx = TestContext::new(vec![]).await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let req = actix_test::TestRequest::get()
|
|
.uri("/readiness")
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
// With no workers, readiness should return SERVICE_UNAVAILABLE
|
|
assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_health_endpoint_details() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![
|
|
MockWorkerConfig {
|
|
port: 18003,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
},
|
|
MockWorkerConfig {
|
|
port: 18004,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
},
|
|
])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let req = actix_test::TestRequest::get().uri("/health").to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
// The health endpoint returns plain text, not JSON
|
|
let body = actix_test::read_body(resp).await;
|
|
let body_str = String::from_utf8_lossy(&body);
|
|
assert!(body_str.contains("All servers healthy"));
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_health_generate_endpoint() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![MockWorkerConfig {
|
|
port: 18005,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let req = actix_test::TestRequest::get()
|
|
.uri("/health_generate")
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
let body: serde_json::Value = actix_test::read_body_json(resp).await;
|
|
assert!(body.is_object());
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod generation_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_generate_success() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![MockWorkerConfig {
|
|
port: 18101,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let payload = json!({
|
|
"text": "Hello, world!",
|
|
"stream": false
|
|
});
|
|
|
|
let req = actix_test::TestRequest::post()
|
|
.uri("/generate")
|
|
.set_json(&payload)
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
let body: serde_json::Value = actix_test::read_body_json(resp).await;
|
|
assert!(body.get("text").is_some());
|
|
assert!(body.get("meta_info").is_some());
|
|
let meta_info = &body["meta_info"];
|
|
assert!(meta_info.get("finish_reason").is_some());
|
|
assert_eq!(meta_info["finish_reason"]["type"], "stop");
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_streaming() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![MockWorkerConfig {
|
|
port: 18102,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let payload = json!({
|
|
"text": "Stream test",
|
|
"stream": true
|
|
});
|
|
|
|
let req = actix_test::TestRequest::post()
|
|
.uri("/generate")
|
|
.set_json(&payload)
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
// Check that it's a streaming response
|
|
let content_type = resp.headers().get("content-type");
|
|
assert!(content_type.is_some());
|
|
assert_eq!(content_type.unwrap(), "text/event-stream");
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_with_worker_failure() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![MockWorkerConfig {
|
|
port: 18103,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 1.0, // Always fail
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let payload = json!({
|
|
"text": "This should fail",
|
|
"stream": false
|
|
});
|
|
|
|
let req = actix_test::TestRequest::post()
|
|
.uri("/generate")
|
|
.set_json(&payload)
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_v1_chat_completions_success() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![MockWorkerConfig {
|
|
port: 18104,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let payload = json!({
|
|
"model": "test-model",
|
|
"messages": [
|
|
{"role": "user", "content": "Hello!"}
|
|
],
|
|
"stream": false
|
|
});
|
|
|
|
let req = actix_test::TestRequest::post()
|
|
.uri("/v1/chat/completions")
|
|
.set_json(&payload)
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
let body: serde_json::Value = actix_test::read_body_json(resp).await;
|
|
assert!(body.get("choices").is_some());
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod model_info_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_get_server_info() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![MockWorkerConfig {
|
|
port: 18201,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let req = actix_test::TestRequest::get()
|
|
.uri("/get_server_info")
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
let body: serde_json::Value = actix_test::read_body_json(resp).await;
|
|
assert!(body.is_object());
|
|
// Check for actual sglang server fields
|
|
assert!(body.get("version").is_some());
|
|
assert!(body.get("model_path").is_some());
|
|
assert!(body.get("tokenizer_path").is_some());
|
|
assert!(body.get("port").is_some());
|
|
assert!(body.get("max_num_batched_tokens").is_some());
|
|
assert!(body.get("schedule_policy").is_some());
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_get_model_info() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![MockWorkerConfig {
|
|
port: 18202,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let req = actix_test::TestRequest::get()
|
|
.uri("/get_model_info")
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
let body: serde_json::Value = actix_test::read_body_json(resp).await;
|
|
assert!(body.is_object());
|
|
// Check for actual sglang model info fields
|
|
assert_eq!(
|
|
body.get("model_path").and_then(|v| v.as_str()),
|
|
Some("mock-model-path")
|
|
);
|
|
assert_eq!(
|
|
body.get("tokenizer_path").and_then(|v| v.as_str()),
|
|
Some("mock-tokenizer-path")
|
|
);
|
|
assert_eq!(
|
|
body.get("is_generation").and_then(|v| v.as_bool()),
|
|
Some(true)
|
|
);
|
|
assert!(body.get("preferred_sampling_params").is_some());
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_v1_models() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![MockWorkerConfig {
|
|
port: 18203,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let req = actix_test::TestRequest::get()
|
|
.uri("/v1/models")
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
let body: serde_json::Value = actix_test::read_body_json(resp).await;
|
|
assert!(body.get("object").is_some());
|
|
assert_eq!(body.get("object").and_then(|v| v.as_str()), Some("list"));
|
|
|
|
let data = body.get("data").and_then(|v| v.as_array());
|
|
assert!(data.is_some());
|
|
|
|
let models = data.unwrap();
|
|
assert!(!models.is_empty());
|
|
|
|
let first_model = &models[0];
|
|
assert_eq!(
|
|
first_model.get("id").and_then(|v| v.as_str()),
|
|
Some("mock-model-v1")
|
|
);
|
|
assert_eq!(
|
|
first_model.get("object").and_then(|v| v.as_str()),
|
|
Some("model")
|
|
);
|
|
assert!(first_model.get("created").is_some());
|
|
assert_eq!(
|
|
first_model.get("owned_by").and_then(|v| v.as_str()),
|
|
Some("sglang")
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_info_with_no_workers() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![]).await;
|
|
let app = ctx.create_app().await;
|
|
|
|
// Test server info with no workers
|
|
let req = actix_test::TestRequest::get()
|
|
.uri("/get_server_info")
|
|
.to_request();
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
// Router may return various error codes when no workers
|
|
assert!(
|
|
resp.status() == StatusCode::OK
|
|
|| resp.status() == StatusCode::SERVICE_UNAVAILABLE
|
|
|| resp.status() == StatusCode::NOT_FOUND
|
|
|| resp.status() == StatusCode::INTERNAL_SERVER_ERROR,
|
|
"Unexpected status code: {:?}",
|
|
resp.status()
|
|
);
|
|
|
|
// Test model info with no workers
|
|
let req = actix_test::TestRequest::get()
|
|
.uri("/get_model_info")
|
|
.to_request();
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
// Router may return various error codes when no workers
|
|
assert!(
|
|
resp.status() == StatusCode::OK
|
|
|| resp.status() == StatusCode::SERVICE_UNAVAILABLE
|
|
|| resp.status() == StatusCode::NOT_FOUND
|
|
|| resp.status() == StatusCode::INTERNAL_SERVER_ERROR,
|
|
"Unexpected status code: {:?}",
|
|
resp.status()
|
|
);
|
|
|
|
// Test v1/models with no workers
|
|
let req = actix_test::TestRequest::get()
|
|
.uri("/v1/models")
|
|
.to_request();
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
// Router may return various error codes when no workers
|
|
assert!(
|
|
resp.status() == StatusCode::OK
|
|
|| resp.status() == StatusCode::SERVICE_UNAVAILABLE
|
|
|| resp.status() == StatusCode::NOT_FOUND
|
|
|| resp.status() == StatusCode::INTERNAL_SERVER_ERROR,
|
|
"Unexpected status code: {:?}",
|
|
resp.status()
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_info_with_multiple_workers() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![
|
|
MockWorkerConfig {
|
|
port: 18204,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
},
|
|
MockWorkerConfig {
|
|
port: 18205,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
},
|
|
])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
// Test that model info is consistent across workers
|
|
for _ in 0..5 {
|
|
let req = actix_test::TestRequest::get()
|
|
.uri("/get_model_info")
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
let body: serde_json::Value = actix_test::read_body_json(resp).await;
|
|
assert_eq!(
|
|
body.get("model_path").and_then(|v| v.as_str()),
|
|
Some("mock-model-path")
|
|
);
|
|
}
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_info_with_unhealthy_worker() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![MockWorkerConfig {
|
|
port: 18206,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 1.0, // Always fail
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let req = actix_test::TestRequest::get()
|
|
.uri("/get_model_info")
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
// Worker with fail_rate: 1.0 should always return an error status
|
|
assert!(
|
|
resp.status() == StatusCode::INTERNAL_SERVER_ERROR
|
|
|| resp.status() == StatusCode::SERVICE_UNAVAILABLE,
|
|
"Expected error status for always-failing worker, got: {:?}",
|
|
resp.status()
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod worker_management_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_add_new_worker() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![]).await;
|
|
let app = ctx.create_app().await;
|
|
|
|
// Start a mock worker
|
|
let mut worker = MockWorker::new(MockWorkerConfig {
|
|
port: 18301,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
});
|
|
let url = worker.start().await.unwrap();
|
|
|
|
// Add the worker
|
|
let req = actix_test::TestRequest::post()
|
|
.uri(&format!("/add_worker?url={}", url))
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
// List workers to verify
|
|
let req = actix_test::TestRequest::get()
|
|
.uri("/list_workers")
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
let body: serde_json::Value = actix_test::read_body_json(resp).await;
|
|
let workers = body["urls"].as_array().unwrap();
|
|
assert!(workers.iter().any(|w| w.as_str().unwrap() == url));
|
|
|
|
worker.stop().await;
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_remove_existing_worker() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![MockWorkerConfig {
|
|
port: 18302,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
// Get the worker URL
|
|
let req = actix_test::TestRequest::get()
|
|
.uri("/list_workers")
|
|
.to_request();
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
let body: serde_json::Value = actix_test::read_body_json(resp).await;
|
|
let workers = body["urls"].as_array().unwrap();
|
|
let worker_url = workers[0].as_str().unwrap();
|
|
|
|
// Remove the worker
|
|
let req = actix_test::TestRequest::post()
|
|
.uri(&format!("/remove_worker?url={}", worker_url))
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
// Verify it's removed
|
|
let req = actix_test::TestRequest::get()
|
|
.uri("/list_workers")
|
|
.to_request();
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
let body: serde_json::Value = actix_test::read_body_json(resp).await;
|
|
let workers = body["urls"].as_array().unwrap();
|
|
assert!(workers.is_empty());
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_add_worker_invalid_url() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![]).await;
|
|
let app = ctx.create_app().await;
|
|
|
|
// Invalid URL format
|
|
let req = actix_test::TestRequest::post()
|
|
.uri("/add_worker?url=not-a-valid-url")
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
|
|
|
|
// Missing URL parameter
|
|
let req = actix_test::TestRequest::post()
|
|
.uri("/add_worker")
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
|
|
|
|
// Empty URL
|
|
let req = actix_test::TestRequest::post()
|
|
.uri("/add_worker?url=")
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_add_duplicate_worker() {
|
|
System::new().block_on(async {
|
|
// Start a mock worker
|
|
let mut worker = MockWorker::new(MockWorkerConfig {
|
|
port: 18303,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
});
|
|
let url = worker.start().await.unwrap();
|
|
|
|
let ctx = TestContext::new(vec![]).await;
|
|
let app = ctx.create_app().await;
|
|
|
|
// Add worker first time
|
|
let req = actix_test::TestRequest::post()
|
|
.uri(&format!("/add_worker?url={}", url))
|
|
.to_request();
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
|
|
|
|
// Try to add same worker again
|
|
let req = actix_test::TestRequest::post()
|
|
.uri(&format!("/add_worker?url={}", url))
|
|
.to_request();
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
// Should return error for duplicate
|
|
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
|
|
|
|
worker.stop().await;
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_add_unhealthy_worker() {
|
|
System::new().block_on(async {
|
|
// Start unhealthy worker
|
|
let mut worker = MockWorker::new(MockWorkerConfig {
|
|
port: 18304,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Unhealthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
});
|
|
let url = worker.start().await.unwrap();
|
|
|
|
let ctx = TestContext::new(vec![]).await;
|
|
let app = ctx.create_app().await;
|
|
|
|
// Try to add unhealthy worker
|
|
let req = actix_test::TestRequest::post()
|
|
.uri(&format!("/add_worker?url={}", url))
|
|
.to_request();
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
|
|
// Router should reject unhealthy workers
|
|
assert!(
|
|
resp.status() == StatusCode::BAD_REQUEST
|
|
|| resp.status() == StatusCode::SERVICE_UNAVAILABLE
|
|
);
|
|
|
|
worker.stop().await;
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod error_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_404_not_found() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![MockWorkerConfig {
|
|
port: 18401,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
// Test unknown endpoint
|
|
let req = actix_test::TestRequest::get()
|
|
.uri("/unknown_endpoint")
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::NOT_FOUND);
|
|
|
|
// Test POST to unknown endpoint
|
|
let req = actix_test::TestRequest::post()
|
|
.uri("/api/v2/generate")
|
|
.set_json(&json!({"text": "test"}))
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::NOT_FOUND);
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_method_not_allowed() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![MockWorkerConfig {
|
|
port: 18402,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
// GET request to POST-only endpoint
|
|
let req = actix_test::TestRequest::get().uri("/generate").to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
// Note: actix-web returns 404 for unmatched methods in some configurations
|
|
assert!(
|
|
resp.status() == StatusCode::METHOD_NOT_ALLOWED
|
|
|| resp.status() == StatusCode::NOT_FOUND
|
|
);
|
|
|
|
// POST request to GET-only endpoint
|
|
let req = actix_test::TestRequest::post()
|
|
.uri("/health")
|
|
.set_json(&json!({}))
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
// Note: actix-web returns 404 for unmatched methods in some configurations
|
|
assert!(
|
|
resp.status() == StatusCode::METHOD_NOT_ALLOWED
|
|
|| resp.status() == StatusCode::NOT_FOUND
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_payload_too_large() {
|
|
System::new().block_on(async {
|
|
// Create context with small payload limit
|
|
let config = RouterConfig {
|
|
mode: RoutingMode::Regular {
|
|
worker_urls: vec![],
|
|
},
|
|
policy: PolicyConfig::Random,
|
|
host: "127.0.0.1".to_string(),
|
|
port: 3010,
|
|
max_payload_size: 1024, // 1KB limit
|
|
request_timeout_secs: 600,
|
|
worker_startup_timeout_secs: 1,
|
|
worker_startup_check_interval_secs: 1,
|
|
discovery: None,
|
|
metrics: None,
|
|
log_dir: None,
|
|
log_level: None,
|
|
};
|
|
|
|
let ctx = TestContext::new_with_config(
|
|
config,
|
|
vec![MockWorkerConfig {
|
|
port: 18403,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}],
|
|
)
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
// Create large payload (> 1KB)
|
|
let large_text = "x".repeat(2000);
|
|
let payload = json!({
|
|
"text": large_text,
|
|
"stream": false
|
|
});
|
|
|
|
let req = actix_test::TestRequest::post()
|
|
.uri("/generate")
|
|
.set_json(&payload)
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
// Note: The test framework may not enforce payload size limits the same way as the full server
|
|
// In production, the server middleware would reject large payloads before reaching handlers
|
|
assert!(
|
|
resp.status() == StatusCode::PAYLOAD_TOO_LARGE || resp.status() == StatusCode::OK
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_invalid_json_payload() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![MockWorkerConfig {
|
|
port: 18404,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
// Send invalid JSON
|
|
let req = actix_test::TestRequest::post()
|
|
.uri("/generate")
|
|
.insert_header(("content-type", "application/json"))
|
|
.set_payload("{invalid json}")
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
|
|
|
|
// Send empty body
|
|
let req = actix_test::TestRequest::post()
|
|
.uri("/generate")
|
|
.insert_header(("content-type", "application/json"))
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_missing_required_fields() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![MockWorkerConfig {
|
|
port: 18405,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
// Missing messages in chat completion
|
|
let payload = json!({
|
|
"model": "test-model"
|
|
// missing "messages"
|
|
});
|
|
|
|
let req = actix_test::TestRequest::post()
|
|
.uri("/v1/chat/completions")
|
|
.set_json(&payload)
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
// Note: Mock worker might accept this, but real implementation would return 400
|
|
// The status depends on the actual router implementation
|
|
assert!(resp.status() == StatusCode::OK || resp.status() == StatusCode::BAD_REQUEST);
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_invalid_model() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![MockWorkerConfig {
|
|
port: 18406,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let payload = json!({
|
|
"model": "invalid-model-name-that-does-not-exist",
|
|
"messages": [{"role": "user", "content": "Hello"}],
|
|
"stream": false
|
|
});
|
|
|
|
let req = actix_test::TestRequest::post()
|
|
.uri("/v1/chat/completions")
|
|
.set_json(&payload)
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
// Mock worker accepts any model, but real implementation might return 400
|
|
assert!(resp.status().is_success() || resp.status() == StatusCode::BAD_REQUEST);
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod cache_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_flush_cache() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![MockWorkerConfig {
|
|
port: 18501,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}])
|
|
.await;
|
|
|
|
let app = actix_test::init_service(
|
|
App::new()
|
|
.app_data(ctx.app_state.clone())
|
|
.service(flush_cache),
|
|
)
|
|
.await;
|
|
|
|
let req = actix_test::TestRequest::post()
|
|
.uri("/flush_cache")
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
// The response might be empty or contain a message
|
|
let body_bytes = actix_test::read_body(resp).await;
|
|
if !body_bytes.is_empty() {
|
|
if let Ok(body) = serde_json::from_slice::<serde_json::Value>(&body_bytes) {
|
|
// Check that we got a successful response with expected fields
|
|
assert!(body.is_object());
|
|
assert!(body.get("message").is_some() || body.get("status").is_some());
|
|
}
|
|
}
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_get_loads() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![
|
|
MockWorkerConfig {
|
|
port: 18502,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
},
|
|
MockWorkerConfig {
|
|
port: 18503,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
},
|
|
])
|
|
.await;
|
|
|
|
let app = actix_test::init_service(
|
|
App::new()
|
|
.app_data(ctx.app_state.clone())
|
|
.service(get_loads),
|
|
)
|
|
.await;
|
|
|
|
let req = actix_test::TestRequest::get()
|
|
.uri("/get_loads")
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
let body: serde_json::Value = actix_test::read_body_json(resp).await;
|
|
|
|
// Verify the response contains load information
|
|
assert!(body.is_object());
|
|
// The exact structure depends on the implementation
|
|
// but should contain worker load information
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_flush_cache_no_workers() {
|
|
System::new().block_on(async {
|
|
let ctx = TestContext::new(vec![]).await;
|
|
|
|
let app = actix_test::init_service(
|
|
App::new()
|
|
.app_data(ctx.app_state.clone())
|
|
.service(flush_cache),
|
|
)
|
|
.await;
|
|
|
|
let req = actix_test::TestRequest::post()
|
|
.uri("/flush_cache")
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
// Should either succeed (no-op) or return service unavailable
|
|
assert!(
|
|
resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod load_balancing_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_request_distribution() {
|
|
System::new().block_on(async {
|
|
// Create multiple workers
|
|
let ctx = TestContext::new(vec![
|
|
MockWorkerConfig {
|
|
port: 18601,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
},
|
|
MockWorkerConfig {
|
|
port: 18602,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
},
|
|
])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
// Send multiple requests and track distribution
|
|
let mut request_count = 0;
|
|
for _ in 0..10 {
|
|
let payload = json!({
|
|
"text": format!("Request {}", request_count),
|
|
"stream": false
|
|
});
|
|
|
|
let req = actix_test::TestRequest::post()
|
|
.uri("/generate")
|
|
.set_json(&payload)
|
|
.to_request();
|
|
|
|
let resp = actix_test::call_service(&app, req).await;
|
|
if resp.status() == StatusCode::OK {
|
|
request_count += 1;
|
|
}
|
|
}
|
|
|
|
// With random policy, all requests should succeed
|
|
assert_eq!(request_count, 10);
|
|
|
|
ctx.shutdown().await;
|
|
});
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod pd_mode_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_pd_mode_routing() {
|
|
System::new().block_on(async {
|
|
// Create PD mode configuration with prefill and decode workers
|
|
let mut prefill_worker = MockWorker::new(MockWorkerConfig {
|
|
port: 18701,
|
|
worker_type: WorkerType::Prefill,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
});
|
|
|
|
let mut decode_worker = MockWorker::new(MockWorkerConfig {
|
|
port: 18702,
|
|
worker_type: WorkerType::Decode,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
});
|
|
|
|
let prefill_url = prefill_worker.start().await.unwrap();
|
|
let decode_url = decode_worker.start().await.unwrap();
|
|
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;
|
|
|
|
// For PD mode, we'll skip the test for now since it requires special handling
|
|
// TODO: Implement PD mode testing with proper worker management
|
|
let _prefill_url = prefill_url;
|
|
let _decode_url = decode_url;
|
|
prefill_worker.stop().await;
|
|
decode_worker.stop().await;
|
|
});
|
|
}
|
|
}
|