310 lines
9.5 KiB
Rust
310 lines
9.5 KiB
Rust
//! Header forwarding integration tests
|
|
//!
|
|
//! Tests for header propagation through the router to workers.
|
|
|
|
mod common;
|
|
|
|
use axum::{
|
|
body::Body,
|
|
extract::Request,
|
|
http::{header::CONTENT_TYPE, StatusCode},
|
|
};
|
|
use common::{
|
|
mock_worker::{HealthStatus, MockWorkerConfig, WorkerType},
|
|
AppTestContext,
|
|
};
|
|
use serde_json::json;
|
|
use smg::config::RouterConfig;
|
|
use tower::ServiceExt;
|
|
|
|
#[cfg(test)]
|
|
mod header_forwarding_tests {
|
|
use super::*;
|
|
|
|
/// Test that X-Request-Id header is forwarded
|
|
#[tokio::test]
|
|
async fn test_request_id_forwarding() {
|
|
let ctx = AppTestContext::new(vec![MockWorkerConfig {
|
|
port: 19400,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let custom_request_id = "test-request-id-12345";
|
|
let payload = json!({
|
|
"text": "Test header forwarding",
|
|
"stream": false
|
|
});
|
|
|
|
let req = Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header(CONTENT_TYPE, "application/json")
|
|
.header("x-request-id", custom_request_id)
|
|
.body(Body::from(serde_json::to_string(&payload).unwrap()))
|
|
.unwrap();
|
|
|
|
let resp = app.oneshot(req).await.unwrap();
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
// Response should have the same request ID
|
|
let response_id = resp.headers().get("x-request-id");
|
|
assert!(response_id.is_some(), "Response should have x-request-id");
|
|
assert_eq!(
|
|
response_id.unwrap().to_str().unwrap(),
|
|
custom_request_id,
|
|
"Request ID should be preserved"
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
}
|
|
|
|
/// Test custom request ID headers
|
|
#[tokio::test]
|
|
async fn test_custom_request_id_headers() {
|
|
let config = RouterConfig::builder()
|
|
.regular_mode(vec![])
|
|
.random_policy()
|
|
.host("127.0.0.1")
|
|
.port(3500)
|
|
.max_payload_size(256 * 1024 * 1024)
|
|
.request_timeout_secs(600)
|
|
.worker_startup_timeout_secs(5)
|
|
.worker_startup_check_interval_secs(1)
|
|
.max_concurrent_requests(64)
|
|
.queue_timeout_secs(60)
|
|
.request_id_headers(vec![
|
|
"custom-trace-id".to_string(),
|
|
"x-correlation-id".to_string(),
|
|
])
|
|
.build_unchecked();
|
|
|
|
let ctx = AppTestContext::new_with_config(
|
|
config,
|
|
vec![MockWorkerConfig {
|
|
port: 19401,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}],
|
|
)
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let custom_trace_id = "my-custom-trace-123";
|
|
let payload = json!({
|
|
"text": "Test custom headers",
|
|
"stream": false
|
|
});
|
|
|
|
let req = Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header(CONTENT_TYPE, "application/json")
|
|
.header("custom-trace-id", custom_trace_id)
|
|
.body(Body::from(serde_json::to_string(&payload).unwrap()))
|
|
.unwrap();
|
|
|
|
let resp = app.oneshot(req).await.unwrap();
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
// Response should use the custom header as request ID
|
|
let response_id = resp.headers().get("x-request-id");
|
|
assert!(response_id.is_some());
|
|
assert_eq!(
|
|
response_id.unwrap().to_str().unwrap(),
|
|
custom_trace_id,
|
|
"Custom header should be used as request ID"
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
}
|
|
|
|
/// Test correlation ID header forwarding
|
|
#[tokio::test]
|
|
async fn test_correlation_id_forwarding() {
|
|
let ctx = AppTestContext::new(vec![MockWorkerConfig {
|
|
port: 19402,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let correlation_id = "correlation-abc-789";
|
|
let payload = json!({
|
|
"text": "Test correlation ID",
|
|
"stream": false
|
|
});
|
|
|
|
let req = Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header(CONTENT_TYPE, "application/json")
|
|
.header("x-correlation-id", correlation_id)
|
|
.body(Body::from(serde_json::to_string(&payload).unwrap()))
|
|
.unwrap();
|
|
|
|
let resp = app.oneshot(req).await.unwrap();
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
// Response should preserve the correlation ID as request ID
|
|
let response_id = resp.headers().get("x-request-id");
|
|
assert!(response_id.is_some());
|
|
assert_eq!(
|
|
response_id.unwrap().to_str().unwrap(),
|
|
correlation_id,
|
|
"Correlation ID should be preserved"
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
}
|
|
|
|
/// Test that request ID is generated when not provided
|
|
#[tokio::test]
|
|
async fn test_auto_generated_request_id() {
|
|
let ctx = AppTestContext::new(vec![MockWorkerConfig {
|
|
port: 19403,
|
|
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": "Test auto-generated ID",
|
|
"stream": false
|
|
});
|
|
|
|
let req = Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header(CONTENT_TYPE, "application/json")
|
|
.body(Body::from(serde_json::to_string(&payload).unwrap()))
|
|
.unwrap();
|
|
|
|
let resp = app.oneshot(req).await.unwrap();
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
// Response should have an auto-generated request ID
|
|
let response_id = resp.headers().get("x-request-id");
|
|
assert!(
|
|
response_id.is_some(),
|
|
"Response should have auto-generated x-request-id"
|
|
);
|
|
|
|
let id_value = response_id.unwrap().to_str().unwrap();
|
|
assert!(!id_value.is_empty(), "Request ID should not be empty");
|
|
// For generate endpoint, ID should have 'gnt-' prefix
|
|
assert!(
|
|
id_value.starts_with("gnt-"),
|
|
"Generate endpoint should have gnt- prefix, got: {}",
|
|
id_value
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
}
|
|
|
|
/// Test chat completions request ID format
|
|
#[tokio::test]
|
|
async fn test_chat_completions_request_id_format() {
|
|
let ctx = AppTestContext::new(vec![MockWorkerConfig {
|
|
port: 19404,
|
|
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 = Request::builder()
|
|
.method("POST")
|
|
.uri("/v1/chat/completions")
|
|
.header(CONTENT_TYPE, "application/json")
|
|
.body(Body::from(serde_json::to_string(&payload).unwrap()))
|
|
.unwrap();
|
|
|
|
let resp = app.oneshot(req).await.unwrap();
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
// Response should have chatcmpl- prefix for chat completions
|
|
let response_id = resp.headers().get("x-request-id");
|
|
assert!(response_id.is_some());
|
|
|
|
let id_value = response_id.unwrap().to_str().unwrap();
|
|
assert!(
|
|
id_value.starts_with("chatcmpl-"),
|
|
"Chat completions should have chatcmpl- prefix, got: {}",
|
|
id_value
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
}
|
|
|
|
/// Test multiple header priorities
|
|
#[tokio::test]
|
|
async fn test_header_priority() {
|
|
let ctx = AppTestContext::new(vec![MockWorkerConfig {
|
|
port: 19405,
|
|
worker_type: WorkerType::Regular,
|
|
health_status: HealthStatus::Healthy,
|
|
response_delay_ms: 0,
|
|
fail_rate: 0.0,
|
|
}])
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let primary_id = "primary-request-id";
|
|
let fallback_id = "fallback-correlation-id";
|
|
let payload = json!({
|
|
"text": "Test header priority",
|
|
"stream": false
|
|
});
|
|
|
|
// When x-request-id is provided, it should take priority
|
|
let req = Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header(CONTENT_TYPE, "application/json")
|
|
.header("x-request-id", primary_id)
|
|
.header("x-correlation-id", fallback_id)
|
|
.body(Body::from(serde_json::to_string(&payload).unwrap()))
|
|
.unwrap();
|
|
|
|
let resp = app.oneshot(req).await.unwrap();
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
let response_id = resp.headers().get("x-request-id");
|
|
assert!(response_id.is_some());
|
|
assert_eq!(
|
|
response_id.unwrap().to_str().unwrap(),
|
|
primary_id,
|
|
"x-request-id should take priority"
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
}
|
|
}
|