[model-gateway] : Rust integration tests for integration_mock replacement (#16441)

.pre-commit-config.yaml

@@ -17,6 +17,7 @@ repos:
       - id: check-merge-conflict
       - id: check-shebang-scripts-are-executable
       - id: detect-private-key
+        exclude: ^sgl-model-gateway/tests/.*_test\.rs$
       - id: debug-statements
       - id: no-commit-to-branch
   - repo: https://github.com/PyCQA/isort

sgl-model-gateway/py_test/unit/__init__.py

@@ -1,7 +0,0 @@
-"""
-Unit tests for sglang_router.
-
-This package contains fast, isolated unit tests for Python components
-of the SGLang router. These tests focus on testing individual functions
-and classes in isolation without starting actual router instances.
-"""

sgl-model-gateway/tests/common/mock_mcp_server.rs

@@ -23,6 +23,12 @@ pub struct MockSearchServer {
     tool_router: ToolRouter<MockSearchServer>,
 }
 
+impl Default for MockSearchServer {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 #[tool_router]
 impl MockSearchServer {
     pub fn new() -> Self {

sgl-model-gateway/tests/common/mod.rs

@@ -6,7 +6,11 @@ pub mod mock_openai_server;
 pub mod mock_worker;
 pub mod streaming_helpers;
 pub mod test_app;
+pub mod test_certs;
+pub mod test_config;
+pub mod tls_mock_worker;
 
+// Re-export commonly used test builders
 use std::{
     fs,
     path::PathBuf,
@@ -33,6 +37,8 @@ use smg::{
     tokenizer::registry::TokenizerRegistry,
     tool_parser::ParserFactory as ToolParserFactory,
 };
+#[allow(unused_imports)]
+pub use test_config::{TestRouterConfig, TestWorkerConfig};
 
 /// Test context for directly testing mock workers without full router setup.
 pub struct WorkerTestContext {

sgl-model-gateway/tests/common/test_certs.rs

@@ -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"
+        );
+    }
+}

sgl-model-gateway/tests/common/test_config.rs

@@ -0,0 +1,316 @@
+//! Test configuration builders to reduce duplication across tests
+//!
+//! Provides pre-configured RouterConfig and MockWorkerConfig builders
+//! for common test scenarios.
+
+use smg::config::{CircuitBreakerConfig, 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 {
+        RouterConfig::builder()
+            .regular_mode(vec![])
+            .policy(PolicyConfig::Manual {
+                eviction_interval_secs: 60,
+                max_idle_secs: 3600,
+            })
+            .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 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);
+    }
+}

sgl-model-gateway/tests/common/tls_mock_worker.rs

@@ -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("/get_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()
+}

sgl-model-gateway/tests/header_forwarding_test.rs

@@ -0,0 +1,309 @@
+//! 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;
+    }
+}

sgl-model-gateway/tests/payload_size_test.rs

@@ -0,0 +1,298 @@
+//! Payload size integration tests
+//!
+//! Tests for request payload size limits and handling.
+
+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 payload_size_tests {
+    use super::*;
+
+    /// Test that small payloads are handled correctly
+    #[tokio::test]
+    async fn test_small_payload() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(4200)
+            .max_payload_size(1024 * 1024) // 1MB limit
+            .request_timeout_secs(600)
+            .worker_startup_timeout_secs(5)
+            .worker_startup_check_interval_secs(1)
+            .max_concurrent_requests(64)
+            .queue_timeout_secs(60)
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![MockWorkerConfig {
+                port: 20200,
+                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": "Small payload test",
+            "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,
+            "Small payload should be accepted"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test that payloads within limit are accepted
+    #[tokio::test]
+    async fn test_payload_within_limit() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(4201)
+            .max_payload_size(1024 * 1024) // 1MB limit
+            .request_timeout_secs(600)
+            .worker_startup_timeout_secs(5)
+            .worker_startup_check_interval_secs(1)
+            .max_concurrent_requests(64)
+            .queue_timeout_secs(60)
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![MockWorkerConfig {
+                port: 20201,
+                worker_type: WorkerType::Regular,
+                health_status: HealthStatus::Healthy,
+                response_delay_ms: 0,
+                fail_rate: 0.0,
+            }],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Create a ~100KB payload (well within 1MB limit)
+        let large_text = "x".repeat(100 * 1024);
+        let payload = json!({
+            "text": large_text,
+            "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,
+            "Payload within limit should be accepted"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test that payloads exceeding limit are rejected
+    #[tokio::test]
+    async fn test_payload_exceeds_limit() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(4202)
+            .max_payload_size(1024) // Very small 1KB limit
+            .request_timeout_secs(600)
+            .worker_startup_timeout_secs(5)
+            .worker_startup_check_interval_secs(1)
+            .max_concurrent_requests(64)
+            .queue_timeout_secs(60)
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![MockWorkerConfig {
+                port: 20202,
+                worker_type: WorkerType::Regular,
+                health_status: HealthStatus::Healthy,
+                response_delay_ms: 0,
+                fail_rate: 0.0,
+            }],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Create a payload larger than 1KB limit
+        let large_text = "x".repeat(2048);
+        let payload = json!({
+            "text": large_text,
+            "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();
+        // Should be rejected with 413 Payload Too Large or similar
+        assert!(
+            resp.status() == StatusCode::PAYLOAD_TOO_LARGE
+                || resp.status() == StatusCode::BAD_REQUEST,
+            "Payload exceeding limit should be rejected, got {}",
+            resp.status()
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test edge case: payload exactly at limit
+    #[tokio::test]
+    async fn test_payload_at_exact_limit() {
+        // Use a more reasonable limit for this test
+        let limit_bytes = 10 * 1024; // 10KB limit
+
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(4203)
+            .max_payload_size(limit_bytes)
+            .request_timeout_secs(600)
+            .worker_startup_timeout_secs(5)
+            .worker_startup_check_interval_secs(1)
+            .max_concurrent_requests(64)
+            .queue_timeout_secs(60)
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![MockWorkerConfig {
+                port: 20203,
+                worker_type: WorkerType::Regular,
+                health_status: HealthStatus::Healthy,
+                response_delay_ms: 0,
+                fail_rate: 0.0,
+            }],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Create a payload slightly under the limit (accounting for JSON overhead)
+        let text_size = limit_bytes - 100; // Leave room for JSON structure
+        let text = "x".repeat(text_size);
+        let payload = json!({
+            "text": text,
+            "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();
+        // Payload at/near limit should be accepted
+        assert!(
+            resp.status() == StatusCode::OK || resp.status() == StatusCode::PAYLOAD_TOO_LARGE,
+            "Payload at limit boundary, got status {}",
+            resp.status()
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test default payload size limit (256MB)
+    #[tokio::test]
+    async fn test_default_payload_limit() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(4204)
+            .max_payload_size(256 * 1024 * 1024) // Default 256MB
+            .request_timeout_secs(600)
+            .worker_startup_timeout_secs(5)
+            .worker_startup_check_interval_secs(1)
+            .max_concurrent_requests(64)
+            .queue_timeout_secs(60)
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![MockWorkerConfig {
+                port: 20204,
+                worker_type: WorkerType::Regular,
+                health_status: HealthStatus::Healthy,
+                response_delay_ms: 0,
+                fail_rate: 0.0,
+            }],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Create a 1MB payload (well within 256MB)
+        let large_text = "x".repeat(1024 * 1024);
+        let payload = json!({
+            "text": large_text,
+            "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,
+            "1MB payload should be accepted with 256MB limit"
+        );
+
+        ctx.shutdown().await;
+    }
+}

sgl-model-gateway/tests/reliability/circuit_breaker_test.rs

@@ -0,0 +1,243 @@
+//! Circuit breaker integration tests
+//!
+//! Tests for circuit breaker behavior: opening, half-open state, recovery, and isolation.
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use serde_json::json;
+use smg::config::{CircuitBreakerConfig, RetryConfig, RouterConfig};
+use tower::ServiceExt;
+
+use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
+
+#[cfg(test)]
+mod circuit_breaker_tests {
+    use super::*;
+
+    /// Test that circuit breaker opens after consecutive failures
+    #[tokio::test]
+    async fn test_circuit_breaker_opens_after_failures() {
+        let config = TestRouterConfig::round_robin_with_circuit_breaker(
+            3200,
+            CircuitBreakerConfig {
+                failure_threshold: 3,
+                success_threshold: 2,
+                timeout_duration_secs: 2,
+                window_duration_secs: 10,
+            },
+        );
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![TestWorkerConfig::flaky(19100, 1.0)], // Always fail
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Make requests until we see failures (500) and then circuit breaker opens (503)
+        let mut saw_500 = false;
+        let mut saw_503 = false;
+
+        for _ in 0..10 {
+            let payload = json!({
+                "text": "Test circuit breaker",
+                "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.clone().oneshot(req).await.unwrap();
+            match resp.status() {
+                StatusCode::INTERNAL_SERVER_ERROR => saw_500 = true,
+                StatusCode::SERVICE_UNAVAILABLE => saw_503 = true,
+                _ => {}
+            }
+        }
+
+        // Should see 500 (worker error) and eventually 503 (circuit breaker open)
+        assert!(
+            saw_500 || saw_503,
+            "Should see either 500 (worker error) or 503 (circuit breaker open)"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test circuit breaker with disabled flag
+    #[tokio::test]
+    async fn test_circuit_breaker_disabled() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(3201)
+            .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)
+            .disable_circuit_breaker()
+            .disable_retries()
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![TestWorkerConfig::flaky(19101, 1.0)], // Always fail
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // With circuit breaker disabled, should always see 500 (never 503)
+        for _ in 0..5 {
+            let payload = json!({
+                "text": "Test disabled CB",
+                "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.clone().oneshot(req).await.unwrap();
+            // With CB disabled, we expect 500 errors (not 503 from CB)
+            assert_eq!(
+                resp.status(),
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Should get 500 with CB disabled, not CB-related 503"
+            );
+        }
+        ctx.shutdown().await;
+    }
+
+    /// Test circuit breaker per-worker isolation
+    #[tokio::test]
+    async fn test_circuit_breaker_per_worker_isolation() {
+        let config = TestRouterConfig::round_robin_with_circuit_breaker(
+            3202,
+            CircuitBreakerConfig {
+                failure_threshold: 2,
+                success_threshold: 1,
+                timeout_duration_secs: 2,
+                window_duration_secs: 10,
+            },
+        );
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(19102, 1.0), // Always fail
+                TestWorkerConfig::healthy(19103),    // Always succeed
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Send enough requests to trigger CB on the failing worker
+        // The healthy worker should continue to serve requests
+        let mut success_count = 0;
+        for _ in 0..20 {
+            let payload = json!({
+                "text": "Test isolation",
+                "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.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        // After CB opens on failing worker, healthy worker should handle all requests
+        // We should see at least some successes
+        assert!(
+            success_count > 0,
+            "Should have some successful requests from healthy worker"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test circuit breaker with retries enabled
+    #[tokio::test]
+    async fn test_circuit_breaker_with_retries() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(3203)
+            .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)
+            .retry_config(RetryConfig {
+                max_retries: 3,
+                initial_backoff_ms: 10,
+                max_backoff_ms: 50,
+                ..Default::default()
+            })
+            .circuit_breaker_config(CircuitBreakerConfig {
+                failure_threshold: 2,
+                success_threshold: 1,
+                timeout_duration_secs: 2,
+                window_duration_secs: 10,
+            })
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(19104, 1.0), // Always fail
+                TestWorkerConfig::healthy(19105),    // Always succeed
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // With retries, requests should succeed by retrying on healthy worker
+        let payload = json!({
+            "text": "Test with retries",
+            "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,
+            "Request should succeed via retry on healthy worker"
+        );
+
+        ctx.shutdown().await;
+    }
+}

sgl-model-gateway/tests/reliability/fault_tolerance_test.rs

@@ -0,0 +1,309 @@
+//! Fault tolerance integration tests
+//!
+//! Tests for system resilience: worker failures, network issues, and recovery.
+
+use std::sync::{
+    atomic::{AtomicUsize, Ordering},
+    Arc,
+};
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use serde_json::json;
+use smg::config::{CircuitBreakerConfig, RetryConfig};
+use tower::ServiceExt;
+
+use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
+
+#[cfg(test)]
+mod fault_tolerance_tests {
+    use super::*;
+
+    /// Test that requests are rerouted when a worker fails
+    #[tokio::test]
+    async fn test_worker_failure_reroute() {
+        let config = TestRouterConfig::round_robin_with_reliability(
+            4100,
+            RetryConfig {
+                max_retries: 3,
+                initial_backoff_ms: 10,
+                max_backoff_ms: 100,
+                ..Default::default()
+            },
+            CircuitBreakerConfig {
+                failure_threshold: 2,
+                success_threshold: 1,
+                timeout_duration_secs: 2,
+                window_duration_secs: 10,
+            },
+        );
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(20100, 1.0), // Always fails
+                TestWorkerConfig::healthy(20101),    // Always succeeds
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Requests should succeed via retry to healthy worker
+        for i in 0..10 {
+            let payload = json!({
+                "text": format!("Fault tolerance test {}", i),
+                "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.clone().oneshot(req).await.unwrap();
+            assert_eq!(
+                resp.status(),
+                StatusCode::OK,
+                "Request should succeed via reroute to healthy worker"
+            );
+        }
+
+        ctx.shutdown().await;
+    }
+
+    /// Test behavior when all workers are temporarily unavailable
+    #[tokio::test]
+    async fn test_all_workers_temporarily_failing() {
+        let config = TestRouterConfig::round_robin_with_retry(
+            4101,
+            RetryConfig {
+                max_retries: 2,
+                initial_backoff_ms: 10,
+                max_backoff_ms: 50,
+                ..Default::default()
+            },
+        );
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(20102, 1.0), // Always fails
+                TestWorkerConfig::flaky(20103, 1.0), // Always fails
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        let payload = json!({
+            "text": "Test with all failing workers",
+            "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();
+        // Should fail when all workers are failing
+        assert!(
+            resp.status() == StatusCode::INTERNAL_SERVER_ERROR
+                || resp.status() == StatusCode::SERVICE_UNAVAILABLE,
+            "Request should fail when all workers are failing, got {}",
+            resp.status()
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test graceful handling of slow workers
+    #[tokio::test]
+    async fn test_slow_worker_handling() {
+        let config = TestRouterConfig::round_robin(4102);
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::slow(20104, 500), // Slow worker
+                TestWorkerConfig::healthy(20105),   // Fast worker
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Send concurrent requests
+        let mut handles = Vec::new();
+        let success_count = Arc::new(AtomicUsize::new(0));
+
+        for i in 0..10 {
+            let app_clone = app.clone();
+            let success_clone = Arc::clone(&success_count);
+
+            let handle = tokio::spawn(async move {
+                let payload = json!({
+                    "text": format!("Slow worker test {}", i),
+                    "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_clone.oneshot(req).await.unwrap();
+                if resp.status() == StatusCode::OK {
+                    success_clone.fetch_add(1, Ordering::SeqCst);
+                }
+            });
+
+            handles.push(handle);
+        }
+
+        for handle in handles {
+            handle.await.unwrap();
+        }
+
+        // All requests should eventually succeed
+        assert_eq!(
+            success_count.load(Ordering::SeqCst),
+            10,
+            "All requests should succeed despite slow worker"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test circuit breaker prevents cascading failures
+    #[tokio::test]
+    async fn test_circuit_breaker_prevents_cascade() {
+        let config = TestRouterConfig::round_robin_with_reliability(
+            4103,
+            RetryConfig {
+                max_retries: 3,
+                initial_backoff_ms: 10,
+                max_backoff_ms: 50,
+                ..Default::default()
+            },
+            CircuitBreakerConfig {
+                failure_threshold: 2,
+                success_threshold: 1,
+                timeout_duration_secs: 5,
+                window_duration_secs: 10,
+            },
+        );
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(20106, 1.0), // Failing worker
+                TestWorkerConfig::healthy(20107),    // Healthy worker
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+        let mut success_count = 0;
+
+        // Send many requests - after CB opens, all should route to healthy worker
+        for i in 0..20 {
+            let payload = json!({
+                "text": format!("Circuit breaker cascade test {}", i),
+                "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.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        // Most requests should succeed after CB opens on failing worker
+        assert!(
+            success_count >= 15,
+            "Most requests should succeed after circuit breaker opens, got {} successes",
+            success_count
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test recovery after worker comes back online (simulated via healthy worker)
+    #[tokio::test]
+    async fn test_system_stability_under_partial_failure() {
+        let config = TestRouterConfig::round_robin_with_reliability(
+            4104,
+            RetryConfig {
+                max_retries: 2,
+                initial_backoff_ms: 10,
+                max_backoff_ms: 50,
+                ..Default::default()
+            },
+            CircuitBreakerConfig {
+                failure_threshold: 3,
+                success_threshold: 1,
+                timeout_duration_secs: 2,
+                window_duration_secs: 10,
+            },
+        );
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(20108, 0.5), // 50% failure rate
+                TestWorkerConfig::healthy(20109),    // Always succeeds
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+        let mut success_count = 0;
+
+        // System should maintain stability with partial failures
+        for i in 0..30 {
+            let payload = json!({
+                "text": format!("Stability test {}", i),
+                "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.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        // With retries and one healthy worker, most should succeed
+        assert!(
+            success_count >= 25,
+            "System should maintain stability under partial failure, got {} successes",
+            success_count
+        );
+
+        ctx.shutdown().await;
+    }
+}

sgl-model-gateway/tests/reliability/mod.rs

@@ -0,0 +1,6 @@
+//! Reliability integration tests
+
+pub mod circuit_breaker_test;
+pub mod fault_tolerance_test;
+pub mod rate_limiting_test;
+pub mod retries_test;

sgl-model-gateway/tests/reliability/rate_limiting_test.rs

@@ -0,0 +1,263 @@
+//! Rate limiting integration tests
+//!
+//! Tests for rate limiting and concurrency control.
+
+use std::sync::{
+    atomic::{AtomicUsize, Ordering},
+    Arc,
+};
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use serde_json::json;
+use smg::config::RouterConfig;
+use tower::ServiceExt;
+
+use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
+
+#[cfg(test)]
+mod rate_limiting_tests {
+    use super::*;
+
+    /// Test that concurrent requests are handled within limits
+    #[tokio::test]
+    async fn test_concurrent_requests_within_limit() {
+        let config = TestRouterConfig::with_concurrency(3400, 10);
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(19300, 50)]).await;
+
+        let app = ctx.create_app().await;
+
+        // Send 5 concurrent requests (within limit of 10)
+        let mut handles = Vec::new();
+        let success_count = Arc::new(AtomicUsize::new(0));
+
+        for i in 0..5 {
+            let app_clone = app.clone();
+            let success_clone = Arc::clone(&success_count);
+
+            let handle = tokio::spawn(async move {
+                let payload = json!({
+                    "text": format!("Concurrent request {}", i),
+                    "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_clone.oneshot(req).await.unwrap();
+                if resp.status() == StatusCode::OK {
+                    success_clone.fetch_add(1, Ordering::SeqCst);
+                }
+            });
+
+            handles.push(handle);
+        }
+
+        // Wait for all requests to complete
+        for handle in handles {
+            handle.await.unwrap();
+        }
+
+        // All requests should succeed within the limit
+        assert_eq!(
+            success_count.load(Ordering::SeqCst),
+            5,
+            "All concurrent requests should succeed within limit"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test rate limit tokens per second
+    #[tokio::test]
+    async fn test_rate_limit_tokens() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .random_policy()
+            .host("127.0.0.1")
+            .port(3401)
+            .max_payload_size(256 * 1024 * 1024)
+            .request_timeout_secs(600)
+            .worker_startup_timeout_secs(5)
+            .worker_startup_check_interval_secs(1)
+            .max_concurrent_requests(100)
+            .rate_limit_tokens_per_second(50) // 50 tokens/sec
+            .queue_timeout_secs(60)
+            .build_unchecked();
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(19301)]).await;
+
+        let app = ctx.create_app().await;
+
+        // Send requests and verify they succeed
+        let mut success_count = 0;
+        for i in 0..10 {
+            let payload = json!({
+                "text": format!("Rate limited request {}", i),
+                "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.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        // Most requests should succeed (some might be rate limited)
+        assert!(
+            success_count >= 5,
+            "At least half of requests should succeed, got {}",
+            success_count
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test unlimited concurrent requests when set to 0
+    #[tokio::test]
+    async fn test_unlimited_concurrent_requests() {
+        let config = TestRouterConfig::with_concurrency(3402, 0); // Unlimited
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(19302, 10)]).await;
+
+        let app = ctx.create_app().await;
+
+        // Send many concurrent requests
+        let mut handles = Vec::new();
+        let success_count = Arc::new(AtomicUsize::new(0));
+
+        for i in 0..20 {
+            let app_clone = app.clone();
+            let success_clone = Arc::clone(&success_count);
+
+            let handle = tokio::spawn(async move {
+                let payload = json!({
+                    "text": format!("Unlimited request {}", i),
+                    "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_clone.oneshot(req).await.unwrap();
+                if resp.status() == StatusCode::OK {
+                    success_clone.fetch_add(1, Ordering::SeqCst);
+                }
+            });
+
+            handles.push(handle);
+        }
+
+        for handle in handles {
+            handle.await.unwrap();
+        }
+
+        // All requests should succeed with unlimited concurrency
+        assert_eq!(
+            success_count.load(Ordering::SeqCst),
+            20,
+            "All requests should succeed with unlimited concurrency"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test queue behavior when requests exceed capacity
+    #[tokio::test]
+    async fn test_queue_behavior() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .random_policy()
+            .host("127.0.0.1")
+            .port(3403)
+            .max_payload_size(256 * 1024 * 1024)
+            .request_timeout_secs(600)
+            .worker_startup_timeout_secs(5)
+            .worker_startup_check_interval_secs(1)
+            .max_concurrent_requests(2) // Very low limit
+            .queue_size(5) // Small queue
+            .queue_timeout_secs(1) // Short timeout
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![TestWorkerConfig::slow(19303, 200)], // Longer delay to test queuing
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Send requests that will exceed capacity
+        let mut handles = Vec::new();
+        let success_count = Arc::new(AtomicUsize::new(0));
+        let error_count = Arc::new(AtomicUsize::new(0));
+
+        for i in 0..10 {
+            let app_clone = app.clone();
+            let success_clone = Arc::clone(&success_count);
+            let error_clone = Arc::clone(&error_count);
+
+            let handle = tokio::spawn(async move {
+                let payload = json!({
+                    "text": format!("Queue test {}", i),
+                    "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_clone.oneshot(req).await.unwrap();
+                if resp.status() == StatusCode::OK {
+                    success_clone.fetch_add(1, Ordering::SeqCst);
+                } else {
+                    error_clone.fetch_add(1, Ordering::SeqCst);
+                }
+            });
+
+            handles.push(handle);
+        }
+
+        for handle in handles {
+            handle.await.unwrap();
+        }
+
+        // Some requests should succeed, some might timeout or be rejected
+        let successes = success_count.load(Ordering::SeqCst);
+        let errors = error_count.load(Ordering::SeqCst);
+
+        assert!(
+            successes > 0,
+            "At least some requests should succeed, got {} successes and {} errors",
+            successes,
+            errors
+        );
+
+        ctx.shutdown().await;
+    }
+}

sgl-model-gateway/tests/reliability/retries_test.rs

@@ -0,0 +1,214 @@
+//! Retry mechanism integration tests
+//!
+//! Tests for retry behavior: exponential backoff, max retries, and retry-on-failure scenarios.
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use serde_json::json;
+use smg::config::{RetryConfig, RouterConfig};
+use tower::ServiceExt;
+
+use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
+
+#[cfg(test)]
+mod retry_tests {
+    use super::*;
+
+    /// Test that retries succeed when at least one worker is healthy
+    #[tokio::test]
+    async fn test_retry_succeeds_with_healthy_fallback() {
+        let retry_config = RetryConfig {
+            max_retries: 3,
+            initial_backoff_ms: 10,
+            max_backoff_ms: 100,
+            ..Default::default()
+        };
+        let config = TestRouterConfig::round_robin_with_retry(3300, retry_config);
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(19200, 1.0), // First worker always fails
+                TestWorkerConfig::healthy(19201),    // Second worker always succeeds
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Request should succeed via retry to healthy worker
+        let payload = json!({
+            "text": "Test retry to healthy worker",
+            "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,
+            "Request should succeed via retry"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test that retries are disabled when configured
+    #[tokio::test]
+    async fn test_retries_disabled() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(3301)
+            .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)
+            .disable_retries()
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![TestWorkerConfig::flaky(19202, 1.0)], // Always fail
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // With retries disabled, request should fail immediately
+        let payload = json!({
+            "text": "Test no retries",
+            "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::INTERNAL_SERVER_ERROR,
+            "Request should fail without retries"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test max retries limit
+    #[tokio::test]
+    async fn test_max_retries_limit() {
+        let retry_config = RetryConfig {
+            max_retries: 2,
+            initial_backoff_ms: 10,
+            max_backoff_ms: 50,
+            ..Default::default()
+        };
+        let config = TestRouterConfig::round_robin_with_retry(3302, retry_config);
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![TestWorkerConfig::flaky(19203, 1.0)], // Always fail
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // All retries will fail, should return error after exhausting retries
+        let payload = json!({
+            "text": "Test max retries",
+            "stream": false
+        });
+
+        let start = std::time::Instant::now();
+        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();
+        let elapsed = start.elapsed();
+
+        // Should eventually fail after retries
+        assert!(
+            resp.status() == StatusCode::INTERNAL_SERVER_ERROR
+                || resp.status() == StatusCode::SERVICE_UNAVAILABLE,
+            "Should fail after exhausting retries, got {}",
+            resp.status()
+        );
+
+        // Should take some time due to backoff (at least initial_backoff_ms)
+        // With 2 retries and 10ms initial backoff, should take at least 10ms
+        // But don't make this too strict as timing can vary
+        assert!(
+            elapsed.as_millis() >= 5,
+            "Should have some backoff delay, got {}ms",
+            elapsed.as_millis()
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test retry with multiple workers - should eventually find healthy one
+    #[tokio::test]
+    async fn test_retry_finds_healthy_worker() {
+        let retry_config = RetryConfig {
+            max_retries: 5,
+            initial_backoff_ms: 5,
+            max_backoff_ms: 50,
+            ..Default::default()
+        };
+        let config = TestRouterConfig::round_robin_with_retry(3303, retry_config);
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(19204, 1.0), // Fail
+                TestWorkerConfig::flaky(19205, 1.0), // Fail
+                TestWorkerConfig::healthy(19206),    // Succeed
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // With round robin and retries, should eventually hit the healthy worker
+        let payload = json!({
+            "text": "Test find healthy worker",
+            "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,
+            "Should succeed by retrying until finding healthy worker"
+        );
+
+        ctx.shutdown().await;
+    }
+}

sgl-model-gateway/tests/reliability_tests.rs

@@ -0,0 +1,8 @@
+//! Reliability integration tests
+
+#[path = "common/mod.rs"]
+pub mod common;
+
+mod reliability;
+
+pub use reliability::*;

sgl-model-gateway/tests/routing/load_balancing_test.rs

@@ -0,0 +1,373 @@
+//! Load balancing integration tests
+//!
+//! Tests for various load balancing policies: round_robin, random, cache_aware, etc.
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use serde_json::json;
+use smg::config::{CircuitBreakerConfig, RetryConfig};
+use tower::ServiceExt;
+
+use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
+
+#[cfg(test)]
+mod round_robin_tests {
+    use super::*;
+
+    /// Test that round robin distributes requests evenly across workers
+    #[tokio::test]
+    async fn test_round_robin_distribution() {
+        let config = TestRouterConfig::round_robin(3100);
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19001, 3))
+                .await;
+
+        let app = ctx.create_app().await;
+        let num_requests = 30;
+        let mut success_count = 0;
+
+        for i in 0..num_requests {
+            let payload = json!({
+                "text": format!("Request {}", i),
+                "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.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        // All requests should succeed with round robin across 3 healthy workers
+        assert_eq!(
+            success_count, num_requests,
+            "All requests should succeed with round robin"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test round robin with one worker failing
+    #[tokio::test]
+    async fn test_round_robin_with_failing_worker() {
+        let config = TestRouterConfig::round_robin_with_retry(
+            3101,
+            RetryConfig {
+                max_retries: 2,
+                initial_backoff_ms: 10,
+                max_backoff_ms: 50,
+                ..Default::default()
+            },
+        );
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(19004, 1.0), // Always fail
+                TestWorkerConfig::healthy(19005),    // Always succeed
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // With retries enabled, requests should eventually succeed
+        // by being retried on the healthy worker
+        let payload = json!({
+            "text": "Test with retries",
+            "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,
+            "Request should succeed via retry on healthy worker"
+        );
+
+        ctx.shutdown().await;
+    }
+}
+
+#[cfg(test)]
+mod random_tests {
+    use super::*;
+
+    /// Test that random policy distributes requests across workers
+    #[tokio::test]
+    async fn test_random_distribution() {
+        let config = TestRouterConfig::random(3102);
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19010, 2))
+                .await;
+
+        let app = ctx.create_app().await;
+        let num_requests = 20;
+        let mut success_count = 0;
+
+        for i in 0..num_requests {
+            let payload = json!({
+                "text": format!("Random request {}", i),
+                "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.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        // All requests should succeed with random policy
+        assert_eq!(
+            success_count, num_requests,
+            "All requests should succeed with random policy"
+        );
+
+        ctx.shutdown().await;
+    }
+}
+
+#[cfg(test)]
+mod cache_aware_tests {
+    use super::*;
+
+    /// Test cache-aware routing uses consistent hashing
+    #[tokio::test]
+    async fn test_cache_aware_consistent_routing() {
+        let config = TestRouterConfig::cache_aware(3103);
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19020, 2))
+                .await;
+
+        let app = ctx.create_app().await;
+
+        // Same prompt should route to same worker (consistent hashing)
+        let same_prompt = "Hello, cache-aware routing test!";
+        let mut worker_ids: Vec<Option<String>> = Vec::new();
+
+        for _ in 0..5 {
+            let payload = json!({
+                "text": same_prompt,
+                "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.clone().oneshot(req).await.unwrap();
+            assert_eq!(resp.status(), StatusCode::OK);
+
+            // Extract worker ID from response header if available
+            let worker_id = resp
+                .headers()
+                .get("x-worker-id")
+                .map(|v| v.to_str().unwrap().to_string());
+            worker_ids.push(worker_id);
+        }
+
+        // All requests should succeed
+        assert_eq!(worker_ids.len(), 5);
+
+        ctx.shutdown().await;
+    }
+
+    /// Test cache-aware routing with different prompts
+    #[tokio::test]
+    async fn test_cache_aware_different_prompts() {
+        let config = TestRouterConfig::cache_aware(3104);
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19022, 2))
+                .await;
+
+        let app = ctx.create_app().await;
+        let prompts = vec![
+            "First unique prompt",
+            "Second unique prompt",
+            "Third unique prompt",
+            "Fourth unique prompt",
+        ];
+
+        let mut success_count = 0;
+        for prompt in prompts {
+            let payload = json!({
+                "text": prompt,
+                "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.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        // All requests should succeed
+        assert_eq!(success_count, 4);
+
+        ctx.shutdown().await;
+    }
+}
+
+#[cfg(test)]
+mod worker_health_tests {
+    use super::*;
+
+    /// Test that failing workers are avoided via circuit breaker/retry
+    #[tokio::test]
+    async fn test_skip_failing_workers() {
+        let config = TestRouterConfig::round_robin_with_reliability(
+            3105,
+            RetryConfig {
+                max_retries: 3,
+                initial_backoff_ms: 10,
+                max_backoff_ms: 50,
+                ..Default::default()
+            },
+            CircuitBreakerConfig {
+                failure_threshold: 2,
+                success_threshold: 1,
+                timeout_duration_secs: 2,
+                window_duration_secs: 10,
+            },
+        );
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(19030, 1.0), // Always fails
+                TestWorkerConfig::healthy(19031),    // Always succeeds
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Requests should succeed by retrying on healthy worker
+        // or by circuit breaker opening on failing worker
+        for i in 0..10 {
+            let payload = json!({
+                "text": format!("Request {}", i),
+                "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.clone().oneshot(req).await.unwrap();
+            assert_eq!(resp.status(), StatusCode::OK);
+        }
+
+        ctx.shutdown().await;
+    }
+
+    /// Test behavior when all workers are unhealthy
+    #[tokio::test]
+    async fn test_all_workers_unhealthy() {
+        let ctx = AppTestContext::new(vec![]).await;
+
+        let app = ctx.create_app().await;
+
+        let payload = json!({
+            "text": "Test with no workers",
+            "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();
+        // Should return service unavailable when no workers available
+        assert!(
+            resp.status() == StatusCode::SERVICE_UNAVAILABLE
+                || resp.status() == StatusCode::INTERNAL_SERVER_ERROR,
+            "Expected 503 or 500 when no workers available, got {}",
+            resp.status()
+        );
+
+        ctx.shutdown().await;
+    }
+}
+
+#[cfg(test)]
+mod worker_response_delay_tests {
+    use super::*;
+
+    /// Test that response delays are handled correctly
+    #[tokio::test]
+    async fn test_worker_with_delay() {
+        let config = TestRouterConfig::random(3106);
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![TestWorkerConfig::slow(19040, 100)], // 100ms delay
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        let start = std::time::Instant::now();
+        let payload = json!({
+            "text": "Test with delay",
+            "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();
+        let elapsed = start.elapsed();
+
+        assert_eq!(resp.status(), StatusCode::OK);
+        // Response should take at least 100ms due to configured delay
+        assert!(
+            elapsed.as_millis() >= 100,
+            "Response should be delayed by at least 100ms, got {}ms",
+            elapsed.as_millis()
+        );
+
+        ctx.shutdown().await;
+    }
+}

sgl-model-gateway/tests/routing/manual_routing_test.rs

@@ -0,0 +1,174 @@
+//! Manual routing policy integration tests
+//!
+//! Tests for the manual routing policy with sticky sessions using X-SMG-Routing-Key header.
+
+use std::collections::{HashMap, HashSet};
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use serde_json::json;
+use tower::ServiceExt;
+
+use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
+
+const ROUTING_KEY_HEADER: &str = "X-SMG-Routing-Key";
+
+#[cfg(test)]
+mod manual_routing_tests {
+    use super::*;
+
+    /// Test sticky routing with X-SMG-Routing-Key header
+    #[tokio::test]
+    async fn test_manual_routing_with_header() {
+        let config = TestRouterConfig::manual(3700);
+
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19700, 2))
+                .await;
+
+        let app = ctx.create_app().await;
+
+        // Send requests with different routing keys
+        let mut key_workers: HashMap<String, HashSet<String>> = HashMap::new();
+
+        for key_id in 0..5 {
+            let routing_key = format!("user-{}", key_id);
+
+            for _ in 0..4 {
+                let payload = json!({
+                    "text": format!("Request for {}", routing_key),
+                    "stream": false
+                });
+
+                let req = Request::builder()
+                    .method("POST")
+                    .uri("/generate")
+                    .header(CONTENT_TYPE, "application/json")
+                    .header(ROUTING_KEY_HEADER, &routing_key)
+                    .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                    .unwrap();
+
+                let resp = app.clone().oneshot(req).await.unwrap();
+                assert_eq!(resp.status(), StatusCode::OK);
+
+                // Extract worker ID from response header
+                if let Some(worker_id) = resp.headers().get("x-worker-id") {
+                    let worker = worker_id.to_str().unwrap().to_string();
+                    key_workers
+                        .entry(routing_key.clone())
+                        .or_default()
+                        .insert(worker);
+                }
+            }
+        }
+
+        // Verify sticky: each routing key should route to exactly one worker
+        for (key, workers) in &key_workers {
+            assert_eq!(
+                workers.len(),
+                1,
+                "Routing key {} should route to exactly one worker, got {:?}",
+                key,
+                workers
+            );
+        }
+
+        ctx.shutdown().await;
+    }
+
+    /// Test random fallback when no routing key header is provided
+    #[tokio::test]
+    async fn test_manual_routing_without_header() {
+        let config = TestRouterConfig::manual(3701);
+
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19702, 2))
+                .await;
+
+        let app = ctx.create_app().await;
+        let mut success_count = 0;
+
+        // Send requests without routing key - should fall back to random selection
+        for i in 0..20 {
+            let payload = json!({
+                "text": format!("Request without key {}", i),
+                "stream": false
+            });
+
+            let req = Request::builder()
+                .method("POST")
+                .uri("/generate")
+                .header(CONTENT_TYPE, "application/json")
+                // No ROUTING_KEY_HEADER
+                .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                .unwrap();
+
+            let resp = app.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        // All requests should succeed via random fallback
+        assert_eq!(
+            success_count, 20,
+            "All requests should succeed with random fallback when no routing key provided"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test that same routing key consistently routes to same worker
+    #[tokio::test]
+    async fn test_manual_routing_consistency() {
+        let config = TestRouterConfig::manual(3702);
+
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19704, 3))
+                .await;
+
+        let app = ctx.create_app().await;
+
+        let routing_key = "consistent-user-123";
+        let mut seen_workers: Vec<String> = Vec::new();
+
+        // Send multiple requests with same routing key
+        for i in 0..10 {
+            let payload = json!({
+                "text": format!("Consistent request {}", i),
+                "stream": false
+            });
+
+            let req = Request::builder()
+                .method("POST")
+                .uri("/generate")
+                .header(CONTENT_TYPE, "application/json")
+                .header(ROUTING_KEY_HEADER, routing_key)
+                .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                .unwrap();
+
+            let resp = app.clone().oneshot(req).await.unwrap();
+            assert_eq!(resp.status(), StatusCode::OK);
+
+            if let Some(worker_id) = resp.headers().get("x-worker-id") {
+                seen_workers.push(worker_id.to_str().unwrap().to_string());
+            }
+        }
+
+        // All requests should go to the same worker
+        if !seen_workers.is_empty() {
+            let first_worker = &seen_workers[0];
+            for worker in &seen_workers {
+                assert_eq!(
+                    worker, first_worker,
+                    "All requests with same routing key should go to same worker"
+                );
+            }
+        }
+
+        ctx.shutdown().await;
+    }
+}

sgl-model-gateway/tests/routing/mod.rs

@@ -0,0 +1,9 @@
+//! Routing integration tests
+
+pub mod cache_aware_backward_compat_test;
+pub mod load_balancing_test;
+pub mod manual_routing_test;
+pub mod pd_routing_test;
+pub mod power_of_two_test;
+pub mod test_openai_routing;
+pub mod test_pd_routing;

sgl-model-gateway/tests/routing/pd_routing_test.rs

@@ -0,0 +1,222 @@
+//! Prefill/Decode (PD) routing integration tests
+//!
+//! Tests for prefill-decode disaggregation routing mode.
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use serde_json::json;
+use smg::config::RouterConfig;
+use tower::ServiceExt;
+
+use crate::common::{
+    mock_worker::{HealthStatus, MockWorkerConfig, WorkerType},
+    AppTestContext, TestWorkerConfig,
+};
+
+#[cfg(test)]
+mod pd_routing_tests {
+    use super::*;
+
+    /// Test basic PD mode routing with prefill and decode workers
+    #[tokio::test]
+    async fn test_pd_mode_basic_routing() {
+        let config = RouterConfig::builder()
+            .prefill_decode_mode(
+                vec![
+                    ("http://127.0.0.1:19800".to_string(), None),
+                    ("http://127.0.0.1:19801".to_string(), None),
+                ],
+                vec![
+                    "http://127.0.0.1:19802".to_string(),
+                    "http://127.0.0.1:19803".to_string(),
+                ],
+            )
+            .power_of_two_policy(1)
+            .host("127.0.0.1")
+            .port(3800)
+            .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)
+            .build_unchecked();
+
+        // Note: For PD mode tests, we need to start prefill and decode workers separately
+        // The test context will need to handle this specially
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                // Prefill workers
+                TestWorkerConfig::prefill(19800),
+                TestWorkerConfig::prefill(19801),
+                // Decode workers
+                TestWorkerConfig::decode(19802),
+                TestWorkerConfig::decode(19803),
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Send requests and verify they succeed
+        for i in 0..10 {
+            let payload = json!({
+                "text": format!("PD mode request {}", i),
+                "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.clone().oneshot(req).await.unwrap();
+            assert_eq!(
+                resp.status(),
+                StatusCode::OK,
+                "PD mode request should succeed"
+            );
+        }
+
+        ctx.shutdown().await;
+    }
+
+    /// Test PD mode with round robin policy
+    #[tokio::test]
+    async fn test_pd_mode_round_robin() {
+        let config = RouterConfig::builder()
+            .prefill_decode_mode(
+                vec![("http://127.0.0.1:19810".to_string(), None)],
+                vec![
+                    "http://127.0.0.1:19811".to_string(),
+                    "http://127.0.0.1:19812".to_string(),
+                ],
+            )
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(3801)
+            .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)
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::prefill(19810),
+                TestWorkerConfig::decode(19811),
+                TestWorkerConfig::decode(19812),
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+        let mut success_count = 0;
+
+        for i in 0..20 {
+            let payload = json!({
+                "text": format!("PD round robin {}", i),
+                "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.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        assert_eq!(
+            success_count, 20,
+            "All requests should succeed in PD mode with round robin"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test PD mode handles worker failures gracefully
+    #[tokio::test]
+    async fn test_pd_mode_with_failing_decode_worker() {
+        use smg::config::RetryConfig;
+
+        let config = RouterConfig::builder()
+            .prefill_decode_mode(
+                vec![("http://127.0.0.1:19820".to_string(), None)],
+                vec![
+                    "http://127.0.0.1:19821".to_string(),
+                    "http://127.0.0.1:19822".to_string(),
+                ],
+            )
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(3802)
+            .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)
+            .retry_config(RetryConfig {
+                max_retries: 3,
+                initial_backoff_ms: 10,
+                max_backoff_ms: 50,
+                ..Default::default()
+            })
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::prefill(19820),
+                MockWorkerConfig {
+                    port: 19821,
+                    worker_type: WorkerType::Decode,
+                    health_status: HealthStatus::Healthy,
+                    response_delay_ms: 0,
+                    fail_rate: 1.0, // Failing decode worker
+                },
+                TestWorkerConfig::decode(19822), // Healthy decode worker
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Request should succeed via retry to healthy decode worker
+        let payload = json!({
+            "text": "Test with failing decode worker",
+            "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,
+            "Request should succeed via retry to healthy decode worker"
+        );
+
+        ctx.shutdown().await;
+    }
+}

sgl-model-gateway/tests/routing/power_of_two_test.rs

@@ -0,0 +1,193 @@
+//! Power of Two load balancing integration tests
+//!
+//! Tests for the Power of Two Choices algorithm that selects the less loaded worker.
+
+use std::sync::{
+    atomic::{AtomicUsize, Ordering},
+    Arc,
+};
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use serde_json::json;
+use smg::config::RouterConfig;
+use tower::ServiceExt;
+
+use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
+
+#[cfg(test)]
+mod power_of_two_tests {
+    use super::*;
+
+    /// Test that power of two distributes requests across workers
+    #[tokio::test]
+    async fn test_power_of_two_distribution() {
+        let config = TestRouterConfig::power_of_two(3600);
+
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19600, 2))
+                .await;
+
+        let app = ctx.create_app().await;
+        let num_requests = 20;
+        let mut success_count = 0;
+
+        for i in 0..num_requests {
+            let payload = json!({
+                "text": format!("Power of two request {}", i),
+                "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.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        // All requests should succeed
+        assert_eq!(
+            success_count, num_requests,
+            "All requests should succeed with power of two policy"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test that power of two prefers less loaded workers
+    #[tokio::test]
+    async fn test_power_of_two_prefers_less_loaded() {
+        let config = TestRouterConfig::power_of_two(3601);
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::slow(19602, 200), // Slow worker
+                TestWorkerConfig::healthy(19603),   // Fast worker
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Send concurrent requests to create load imbalance
+        let mut handles = Vec::new();
+        let success_count = Arc::new(AtomicUsize::new(0));
+
+        for i in 0..30 {
+            let app_clone = app.clone();
+            let success_clone = Arc::clone(&success_count);
+
+            let handle = tokio::spawn(async move {
+                let payload = json!({
+                    "text": format!("Concurrent request {}", i),
+                    "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_clone.oneshot(req).await.unwrap();
+                if resp.status() == StatusCode::OK {
+                    success_clone.fetch_add(1, Ordering::SeqCst);
+                }
+            });
+
+            handles.push(handle);
+        }
+
+        for handle in handles {
+            handle.await.unwrap();
+        }
+
+        // All requests should succeed
+        assert!(
+            success_count.load(Ordering::SeqCst) >= 25,
+            "Most requests should succeed with power of two"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test power of two with failing worker uses retry/CB to route to healthy worker
+    #[tokio::test]
+    async fn test_power_of_two_with_failing_worker() {
+        use smg::config::{CircuitBreakerConfig, RetryConfig};
+
+        let retry_config = RetryConfig {
+            max_retries: 3,
+            initial_backoff_ms: 10,
+            max_backoff_ms: 50,
+            ..Default::default()
+        };
+        let circuit_breaker = CircuitBreakerConfig {
+            failure_threshold: 2,
+            success_threshold: 1,
+            timeout_duration_secs: 2,
+            window_duration_secs: 10,
+        };
+
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .power_of_two_policy(1)
+            .host("127.0.0.1")
+            .port(3602)
+            .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)
+            .retry_config(retry_config)
+            .circuit_breaker_config(circuit_breaker)
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(19604, 1.0), // Always fails
+                TestWorkerConfig::healthy(19605),    // Always succeeds
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Requests should succeed via retry to healthy worker
+        for i in 0..10 {
+            let payload = json!({
+                "text": format!("Request with failing worker {}", i),
+                "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.clone().oneshot(req).await.unwrap();
+            assert_eq!(
+                resp.status(),
+                StatusCode::OK,
+                "Request should succeed via retry on healthy worker"
+            );
+        }
+
+        ctx.shutdown().await;
+    }
+}

sgl-model-gateway/tests/routing/test_openai_routing.rs

@@ -37,8 +37,7 @@ use tokio::{
 };
 use tower::ServiceExt;
 
-mod common;
-use common::mock_openai_server::MockOpenAIServer;
+use crate::common::mock_openai_server::MockOpenAIServer;
 
 /// Helper function to create a minimal chat completion request for testing
 fn create_minimal_chat_request() -> ChatCompletionRequest {
@@ -95,9 +94,9 @@ fn create_minimal_completion_request() -> CompletionRequest {
 /// Test basic OpenAI router creation and configuration
 #[tokio::test]
 async fn test_openai_router_creation() {
-    let ctx = common::test_app::create_test_app_context().await;
+    let ctx = crate::common::test_app::create_test_app_context().await;
     // Register an external worker before creating the router
-    common::test_app::register_external_worker(&ctx, "https://api.openai.com", None);
+    crate::common::test_app::register_external_worker(&ctx, "https://api.openai.com", None);
     let router = OpenAIRouter::new(&ctx).await;
 
     assert!(router.is_ok(), "Router creation should succeed");
@@ -110,8 +109,8 @@ async fn test_openai_router_creation() {
 /// Test server info endpoint
 #[tokio::test]
 async fn test_openai_router_server_info() {
-    let ctx = common::test_app::create_test_app_context().await;
-    common::test_app::register_external_worker(&ctx, "https://api.openai.com", None);
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    crate::common::test_app::register_external_worker(&ctx, "https://api.openai.com", None);
     let router = OpenAIRouter::new(&ctx).await.unwrap();
 
     let req = Request::builder()
@@ -135,8 +134,8 @@ async fn test_openai_router_server_info() {
 async fn test_openai_router_models() {
     // Use mock server for deterministic models response
     let mock_server = MockOpenAIServer::new().await;
-    let ctx = common::test_app::create_test_app_context().await;
-    common::test_app::register_external_worker(&ctx, &mock_server.base_url(), None);
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    crate::common::test_app::register_external_worker(&ctx, &mock_server.base_url(), None);
     let router = OpenAIRouter::new(&ctx).await.unwrap();
 
     let req = Request::builder()
@@ -208,8 +207,8 @@ async fn test_openai_router_responses_with_mock() {
 
     let base_url = format!("http://{}", addr);
 
-    let ctx = common::test_app::create_test_app_context().await;
-    common::test_app::register_external_worker(&ctx, &base_url, Some(vec!["gpt-4o-mini"]));
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    crate::common::test_app::register_external_worker(&ctx, &base_url, Some(vec!["gpt-4o-mini"]));
     let router = OpenAIRouter::new(&ctx).await.unwrap();
 
     // Get storage from context (router uses this, not a separate storage)
@@ -473,8 +472,8 @@ async fn test_openai_router_responses_streaming_with_mock() {
 
     let base_url = format!("http://{}", addr);
 
-    let ctx = common::test_app::create_test_app_context().await;
-    common::test_app::register_external_worker(&ctx, &base_url, Some(vec!["gpt-5-nano"]));
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    crate::common::test_app::register_external_worker(&ctx, &base_url, Some(vec!["gpt-5-nano"]));
     let router = OpenAIRouter::new(&ctx).await.unwrap();
 
     // Get storage from context and seed a previous response
@@ -581,7 +580,7 @@ async fn test_router_factory_openai_mode() {
 
     let router_config = RouterConfig::new(routing_mode, smg::config::PolicyConfig::Random);
 
-    let app_context = common::create_test_context(router_config).await;
+    let app_context = crate::common::create_test_context(router_config).await;
 
     let router = smg::routers::RouterFactory::create_router(&app_context).await;
     assert!(
@@ -596,8 +595,8 @@ async fn test_router_factory_openai_mode() {
 /// Test that unsupported endpoints return proper error codes
 #[tokio::test]
 async fn test_unsupported_endpoints() {
-    let ctx = common::test_app::create_test_app_context().await;
-    common::test_app::register_external_worker(&ctx, "https://api.openai.com", None);
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    crate::common::test_app::register_external_worker(&ctx, "https://api.openai.com", None);
     let router = OpenAIRouter::new(&ctx).await.unwrap();
 
     let generate_request = GenerateRequest {
@@ -655,9 +654,9 @@ async fn test_openai_router_chat_completion_with_mock() {
     let mock_server = MockOpenAIServer::new().await;
     let base_url = mock_server.base_url();
 
-    let ctx = common::test_app::create_test_app_context().await;
+    let ctx = crate::common::test_app::create_test_app_context().await;
     // Register the mock server worker and create router
-    common::test_app::register_external_worker(&ctx, &base_url, None);
+    crate::common::test_app::register_external_worker(&ctx, &base_url, None);
     let router = OpenAIRouter::new(&ctx).await.unwrap();
 
     // Create a minimal chat completion request
@@ -691,9 +690,9 @@ async fn test_openai_e2e_with_server() {
     let mock_server = MockOpenAIServer::new().await;
     let base_url = mock_server.base_url();
 
-    let ctx = common::test_app::create_test_app_context().await;
+    let ctx = crate::common::test_app::create_test_app_context().await;
     // Register the mock server worker and create router
-    common::test_app::register_external_worker(&ctx, &base_url, None);
+    crate::common::test_app::register_external_worker(&ctx, &base_url, None);
     let router = OpenAIRouter::new(&ctx).await.unwrap();
 
     // Create Axum app with chat completions endpoint
@@ -757,8 +756,8 @@ async fn test_openai_e2e_with_server() {
 async fn test_openai_router_chat_streaming_with_mock() {
     let mock_server = MockOpenAIServer::new().await;
     let base_url = mock_server.base_url();
-    let ctx = common::test_app::create_test_app_context().await;
-    common::test_app::register_external_worker(&ctx, &base_url, None);
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    crate::common::test_app::register_external_worker(&ctx, &base_url, None);
     let router = OpenAIRouter::new(&ctx).await.unwrap();
 
     // Build a streaming chat request
@@ -797,8 +796,12 @@ async fn test_openai_router_chat_streaming_with_mock() {
 /// Test circuit breaker functionality
 #[tokio::test]
 async fn test_openai_router_circuit_breaker() {
-    let ctx = common::test_app::create_test_app_context().await;
-    common::test_app::register_external_worker(&ctx, "http://invalid-url-that-will-fail", None);
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    crate::common::test_app::register_external_worker(
+        &ctx,
+        "http://invalid-url-that-will-fail",
+        None,
+    );
     let router = OpenAIRouter::new(&ctx).await.unwrap();
 
     let chat_request = create_minimal_chat_request();
@@ -821,9 +824,9 @@ async fn test_openai_router_circuit_breaker() {
 /// requests to workers, not for the models endpoint.
 #[tokio::test]
 async fn test_openai_router_models_from_registry() {
-    let ctx = common::test_app::create_test_app_context().await;
+    let ctx = crate::common::test_app::create_test_app_context().await;
     // Register a worker with the default model
-    common::test_app::register_external_worker(&ctx, "https://api.example.com", None);
+    crate::common::test_app::register_external_worker(&ctx, "https://api.example.com", None);
     let router = OpenAIRouter::new(&ctx).await.unwrap();
 
     // Get models - should return the registered model

sgl-model-gateway/tests/routing/test_pd_routing.rs

@@ -1,5 +1,5 @@
 #[cfg(test)]
-mod test_pd_routing {
+mod pd_routing_unit_tests {
     use serde_json::json;
     use smg::{
         app_context::AppContext,

sgl-model-gateway/tests/routing_tests.rs

@@ -0,0 +1,8 @@
+//! Routing integration tests
+
+#[path = "common/mod.rs"]
+pub mod common;
+
+mod routing;
+
+pub use routing::*;

sgl-model-gateway/tests/security/auth_test.rs

@@ -0,0 +1,243 @@
+//! Authentication and authorization integration tests
+//!
+//! Tests for API key enforcement and access control.
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use serde_json::json;
+use tower::ServiceExt;
+
+use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
+
+const AUTH_HEADER: &str = "Authorization";
+
+#[cfg(test)]
+mod auth_tests {
+    use super::*;
+
+    /// Test request without API key when auth is not required
+    #[tokio::test]
+    async fn test_no_auth_required() {
+        let config = TestRouterConfig::round_robin(4300);
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20300)]).await;
+
+        let app = ctx.create_app().await;
+
+        // Request without auth header should succeed when no auth required
+        let payload = json!({
+            "text": "Test without auth",
+            "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,
+            "Request without auth should succeed when no auth required"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test request with valid API key format
+    #[tokio::test]
+    async fn test_with_api_key_header() {
+        let config = TestRouterConfig::round_robin(4301);
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20301)]).await;
+
+        let app = ctx.create_app().await;
+
+        // Request with Bearer token header
+        let payload = json!({
+            "text": "Test with auth header",
+            "stream": false
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/generate")
+            .header(CONTENT_TYPE, "application/json")
+            .header(AUTH_HEADER, "Bearer test-api-key-12345")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        // Without auth enforcement, request should succeed
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "Request with auth header should succeed"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test health endpoint doesn't require authentication
+    #[tokio::test]
+    async fn test_health_endpoint_no_auth() {
+        let config = TestRouterConfig::round_robin(4302);
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20302)]).await;
+
+        let app = ctx.create_app().await;
+
+        // Health endpoint should be accessible without auth
+        let req = Request::builder()
+            .method("GET")
+            .uri("/health")
+            .body(Body::empty())
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "Health endpoint should not require auth"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test OpenAI-compatible API key header (X-API-Key)
+    #[tokio::test]
+    async fn test_openai_api_key_header() {
+        let config = TestRouterConfig::round_robin(4303);
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20303)]).await;
+
+        let app = ctx.create_app().await;
+
+        // Request with X-API-Key header (OpenAI style)
+        let payload = json!({
+            "text": "Test with X-API-Key",
+            "stream": false
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/generate")
+            .header(CONTENT_TYPE, "application/json")
+            .header("X-API-Key", "sk-test-key-12345")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "Request with X-API-Key should succeed"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test multiple concurrent authenticated requests
+    #[tokio::test]
+    async fn test_concurrent_authenticated_requests() {
+        use std::sync::{
+            atomic::{AtomicUsize, Ordering},
+            Arc,
+        };
+
+        let config = TestRouterConfig::round_robin(4304);
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20304)]).await;
+
+        let app = ctx.create_app().await;
+        let success_count = Arc::new(AtomicUsize::new(0));
+        let mut handles = Vec::new();
+
+        for i in 0..20 {
+            let app_clone = app.clone();
+            let success_clone = Arc::clone(&success_count);
+
+            let handle = tokio::spawn(async move {
+                let payload = json!({
+                    "text": format!("Concurrent auth test {}", i),
+                    "stream": false
+                });
+
+                let req = Request::builder()
+                    .method("POST")
+                    .uri("/generate")
+                    .header(CONTENT_TYPE, "application/json")
+                    .header(AUTH_HEADER, format!("Bearer test-key-{}", i))
+                    .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                    .unwrap();
+
+                let resp = app_clone.oneshot(req).await.unwrap();
+                if resp.status() == StatusCode::OK {
+                    success_clone.fetch_add(1, Ordering::SeqCst);
+                }
+            });
+
+            handles.push(handle);
+        }
+
+        for handle in handles {
+            handle.await.unwrap();
+        }
+
+        assert_eq!(
+            success_count.load(Ordering::SeqCst),
+            20,
+            "All concurrent authenticated requests should succeed"
+        );
+
+        ctx.shutdown().await;
+    }
+}
+
+#[cfg(test)]
+mod mtls_tests {
+    use super::*;
+
+    /// Test that TLS configuration options exist
+    /// Note: Actual mTLS testing would require certificate setup
+    #[tokio::test]
+    async fn test_tls_config_available() {
+        // This test verifies the config builder accepts TLS-related options
+        // Actual mTLS testing requires certificate infrastructure
+        let config = TestRouterConfig::round_robin(4305);
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20305)]).await;
+
+        let app = ctx.create_app().await;
+
+        // Basic request should work (no TLS in test mode)
+        let payload = json!({
+            "text": "TLS config test",
+            "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);
+
+        ctx.shutdown().await;
+    }
+}

sgl-model-gateway/tests/security/mod.rs

@@ -0,0 +1,5 @@
+//! Security integration tests
+
+pub mod auth_integration_test;
+pub mod auth_test;
+pub mod mtls_test;

sgl-model-gateway/tests/security/mtls_test.rs

@@ -0,0 +1,506 @@
+//! mTLS (Mutual TLS) integration tests
+//!
+//! Tests for TLS and mTLS communication between router and workers.
+//! Covers:
+//! - Successful mTLS communication with client certificates
+//! - TLS failure without client certificate when required
+//! - TLS-only mode (server authentication only)
+//! - TLS failure without CA certificate
+
+use std::{io::BufReader, time::Duration};
+
+use rustls::{ClientConfig, RootCertStore};
+use rustls_pemfile::{certs, pkcs8_private_keys};
+use serde_json::json;
+
+use crate::common::{
+    test_certs::TestCertificates,
+    tls_mock_worker::{TlsMockWorker, TlsMockWorkerConfig},
+};
+
+// Ensure crypto provider is installed
+fn ensure_crypto_provider() {
+    use std::sync::Once;
+    static INIT: Once = Once::new();
+    INIT.call_once(|| {
+        let _ = rustls::crypto::ring::default_provider().install_default();
+    });
+}
+
+/// Helper to create a TLS-enabled reqwest client with client certificates
+fn create_mtls_client(
+    test_certs: &TestCertificates,
+) -> Result<reqwest::Client, Box<dyn std::error::Error>> {
+    ensure_crypto_provider();
+
+    // Read CA certificate
+    let ca_file = std::fs::File::open(&test_certs.ca_cert_path)?;
+    let mut ca_reader = BufReader::new(ca_file);
+    let ca_certs: Vec<_> = certs(&mut ca_reader).filter_map(|r| r.ok()).collect();
+
+    let mut root_store = RootCertStore::empty();
+    for cert in ca_certs {
+        root_store.add(cert)?;
+    }
+
+    // Read client certificate
+    let cert_file = std::fs::File::open(&test_certs.client_cert_path)?;
+    let mut cert_reader = BufReader::new(cert_file);
+    let client_certs: Vec<_> = certs(&mut cert_reader).filter_map(|r| r.ok()).collect();
+
+    // Read client key
+    let key_file = std::fs::File::open(&test_certs.client_key_path)?;
+    let mut key_reader = BufReader::new(key_file);
+    let client_key = pkcs8_private_keys(&mut key_reader)
+        .next()
+        .ok_or("No private key found")??;
+
+    // Build client config with client certificate
+    let client_config = ClientConfig::builder()
+        .with_root_certificates(root_store)
+        .with_client_auth_cert(client_certs, client_key.into())
+        .map_err(|e| format!("Failed to build client config: {}", e))?;
+
+    let client = reqwest::Client::builder()
+        .use_preconfigured_tls(client_config)
+        .timeout(Duration::from_secs(10))
+        .build()?;
+
+    Ok(client)
+}
+
+/// Helper to create a TLS client without client certificates (for server-auth-only)
+fn create_tls_client(
+    test_certs: &TestCertificates,
+) -> Result<reqwest::Client, Box<dyn std::error::Error>> {
+    ensure_crypto_provider();
+
+    // Read CA certificate
+    let ca_file = std::fs::File::open(&test_certs.ca_cert_path)?;
+    let mut ca_reader = BufReader::new(ca_file);
+    let ca_certs: Vec<_> = certs(&mut ca_reader).filter_map(|r| r.ok()).collect();
+
+    let mut root_store = RootCertStore::empty();
+    for cert in ca_certs {
+        root_store.add(cert)?;
+    }
+
+    // Build client config without client certificate
+    let client_config = ClientConfig::builder()
+        .with_root_certificates(root_store)
+        .with_no_client_auth();
+
+    let client = reqwest::Client::builder()
+        .use_preconfigured_tls(client_config)
+        .timeout(Duration::from_secs(10))
+        .build()?;
+
+    Ok(client)
+}
+
+/// Helper to create a client without CA certificate (for failure test)
+fn create_client_without_ca() -> Result<reqwest::Client, Box<dyn std::error::Error>> {
+    ensure_crypto_provider();
+
+    // Build client with empty root store (will fail to verify server)
+    let root_store = RootCertStore::empty();
+    let client_config = ClientConfig::builder()
+        .with_root_certificates(root_store)
+        .with_no_client_auth();
+
+    let client = reqwest::Client::builder()
+        .use_preconfigured_tls(client_config)
+        .timeout(Duration::from_secs(10))
+        .build()?;
+
+    Ok(client)
+}
+
+#[cfg(test)]
+mod mtls_tests {
+    use super::*;
+
+    /// Test successful mTLS communication between client and TLS worker
+    ///
+    /// This test verifies that:
+    /// 1. TLS mock worker starts with mTLS configuration
+    /// 2. Client with proper certificates can connect and communicate
+    /// 3. Requests succeed with proper authentication
+    #[tokio::test]
+    async fn test_mtls_successful_communication() {
+        // Generate test certificates
+        let certs = TestCertificates::generate().expect("Failed to generate test certificates");
+
+        // Start mTLS-enabled mock worker
+        let mut worker = TlsMockWorker::new(TlsMockWorkerConfig {
+            port: 0, // Auto-assign port
+            require_client_cert: true,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        });
+
+        let worker_url = worker
+            .start(
+                &certs.server_cert_path,
+                &certs.server_key_path,
+                Some(&certs.ca_cert_path),
+            )
+            .await
+            .expect("Failed to start mTLS worker");
+
+        // Create client with full mTLS credentials
+        let client = create_mtls_client(&certs).expect("Failed to create mTLS client");
+
+        // Test health endpoint
+        let health_resp = client
+            .get(format!("{}/health", worker_url))
+            .send()
+            .await
+            .expect("Health request failed");
+
+        assert!(
+            health_resp.status().is_success(),
+            "Health check should succeed with mTLS: {}",
+            health_resp.status()
+        );
+
+        let health_json: serde_json::Value = health_resp.json().await.unwrap();
+        assert_eq!(health_json["status"], "healthy");
+        assert_eq!(health_json["tls_enabled"], true);
+
+        // Test generate endpoint
+        let payload = json!({
+            "text": "Test mTLS request",
+            "stream": false
+        });
+
+        let gen_resp = client
+            .post(format!("{}/generate", worker_url))
+            .json(&payload)
+            .send()
+            .await
+            .expect("Generate request failed");
+
+        assert!(
+            gen_resp.status().is_success(),
+            "Generate should succeed with mTLS: {}",
+            gen_resp.status()
+        );
+
+        let gen_json: serde_json::Value = gen_resp.json().await.unwrap();
+        assert!(gen_json["text"].as_str().is_some());
+        assert_eq!(gen_json["meta_info"]["tls_verified"], true);
+
+        // Cleanup
+        worker.stop().await;
+    }
+
+    /// Test that mTLS worker rejects connections without client certificate
+    ///
+    /// This test verifies that:
+    /// 1. mTLS worker requires client certificate
+    /// 2. Connection without client cert fails
+    #[tokio::test]
+    async fn test_mtls_failure_without_client_cert() {
+        // Generate test certificates
+        let certs = TestCertificates::generate().expect("Failed to generate test certificates");
+
+        // Start mTLS-enabled mock worker (requires client cert)
+        let mut worker = TlsMockWorker::new(TlsMockWorkerConfig {
+            port: 0,
+            require_client_cert: true,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        });
+
+        let worker_url = worker
+            .start(
+                &certs.server_cert_path,
+                &certs.server_key_path,
+                Some(&certs.ca_cert_path),
+            )
+            .await
+            .expect("Failed to start mTLS worker");
+
+        // Create client WITHOUT client certificate (only has CA for server verification)
+        let client = create_tls_client(&certs).expect("Failed to create TLS-only client");
+
+        // Attempt to connect - should fail because no client cert provided
+        let result = client.get(format!("{}/health", worker_url)).send().await;
+
+        // Connection should fail or be rejected
+        assert!(
+            result.is_err(),
+            "Connection should fail without client certificate"
+        );
+
+        // Cleanup
+        worker.stop().await;
+    }
+
+    /// Test TLS-only mode (server authentication only, no client cert required)
+    ///
+    /// This test verifies that:
+    /// 1. TLS worker can operate without requiring client certificates
+    /// 2. Client can connect with just CA certificate for server verification
+    #[tokio::test]
+    async fn test_tls_server_auth_only() {
+        // Generate test certificates
+        let certs = TestCertificates::generate().expect("Failed to generate test certificates");
+
+        // Start TLS-only mock worker (does NOT require client cert)
+        let mut worker = TlsMockWorker::new(TlsMockWorkerConfig {
+            port: 0,
+            require_client_cert: false, // TLS-only mode
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        });
+
+        let worker_url = worker
+            .start(
+                &certs.server_cert_path,
+                &certs.server_key_path,
+                None, // No CA needed for client verification
+            )
+            .await
+            .expect("Failed to start TLS worker");
+
+        // Create client with just CA cert for server verification (no client cert)
+        let client = create_tls_client(&certs).expect("Failed to create TLS client");
+
+        // Test health endpoint - should succeed without client cert
+        let health_resp = client
+            .get(format!("{}/health", worker_url))
+            .send()
+            .await
+            .expect("Health request failed");
+
+        assert!(
+            health_resp.status().is_success(),
+            "Health check should succeed in TLS-only mode: {}",
+            health_resp.status()
+        );
+
+        let health_json: serde_json::Value = health_resp.json().await.unwrap();
+        assert_eq!(health_json["status"], "healthy");
+
+        // Test chat completions endpoint
+        let payload = json!({
+            "model": "mock-tls-model",
+            "messages": [{"role": "user", "content": "Hello TLS"}]
+        });
+
+        let chat_resp = client
+            .post(format!("{}/v1/chat/completions", worker_url))
+            .json(&payload)
+            .send()
+            .await
+            .expect("Chat request failed");
+
+        assert!(
+            chat_resp.status().is_success(),
+            "Chat should succeed in TLS-only mode: {}",
+            chat_resp.status()
+        );
+
+        // Cleanup
+        worker.stop().await;
+    }
+
+    /// Test TLS failure when client doesn't have CA certificate
+    ///
+    /// This test verifies that:
+    /// 1. Client cannot verify server without proper CA certificate
+    /// 2. Connection fails due to certificate verification
+    #[tokio::test]
+    async fn test_tls_failure_without_ca_cert() {
+        // Generate test certificates
+        let certs = TestCertificates::generate().expect("Failed to generate test certificates");
+
+        // Start TLS mock worker
+        let mut worker = TlsMockWorker::new(TlsMockWorkerConfig {
+            port: 0,
+            require_client_cert: false,
+            response_delay_ms: 0,
+            fail_rate: 0.0,
+        });
+
+        let worker_url = worker
+            .start(&certs.server_cert_path, &certs.server_key_path, None)
+            .await
+            .expect("Failed to start TLS worker");
+
+        // Create client WITHOUT CA certificate
+        let client = create_client_without_ca().expect("Failed to create client without CA");
+
+        // Attempt to connect - should fail because cannot verify server cert
+        let result = client.get(format!("{}/health", worker_url)).send().await;
+
+        // Connection should fail due to certificate verification
+        assert!(
+            result.is_err(),
+            "Connection should fail without CA certificate for server verification"
+        );
+
+        // Cleanup
+        worker.stop().await;
+    }
+
+    /// Test multiple concurrent mTLS requests
+    ///
+    /// This test verifies that mTLS connections work correctly under concurrent load
+    #[tokio::test]
+    async fn test_mtls_concurrent_requests() {
+        use std::sync::{
+            atomic::{AtomicUsize, Ordering},
+            Arc,
+        };
+
+        // Generate test certificates
+        let certs = TestCertificates::generate().expect("Failed to generate test certificates");
+
+        // Start mTLS-enabled mock worker
+        let mut worker = TlsMockWorker::new(TlsMockWorkerConfig {
+            port: 0,
+            require_client_cert: true,
+            response_delay_ms: 10, // Small delay to simulate work
+            fail_rate: 0.0,
+        });
+
+        let worker_url = worker
+            .start(
+                &certs.server_cert_path,
+                &certs.server_key_path,
+                Some(&certs.ca_cert_path),
+            )
+            .await
+            .expect("Failed to start mTLS worker");
+
+        // Create mTLS client
+        let client = Arc::new(create_mtls_client(&certs).expect("Failed to create mTLS client"));
+        let success_count = Arc::new(AtomicUsize::new(0));
+        let worker_url = Arc::new(worker_url);
+
+        let mut handles = Vec::new();
+
+        // Spawn concurrent requests
+        for i in 0..10 {
+            let client_clone = Arc::clone(&client);
+            let success_clone = Arc::clone(&success_count);
+            let url_clone = Arc::clone(&worker_url);
+
+            let handle = tokio::spawn(async move {
+                let payload = json!({
+                    "text": format!("Concurrent mTLS request {}", i),
+                    "stream": false
+                });
+
+                let resp = client_clone
+                    .post(format!("{}/generate", url_clone))
+                    .json(&payload)
+                    .send()
+                    .await;
+
+                if let Ok(response) = resp {
+                    if response.status().is_success() {
+                        success_clone.fetch_add(1, Ordering::SeqCst);
+                    }
+                }
+            });
+
+            handles.push(handle);
+        }
+
+        // Wait for all requests to complete
+        for handle in handles {
+            handle.await.unwrap();
+        }
+
+        // All requests should succeed
+        assert_eq!(
+            success_count.load(Ordering::SeqCst),
+            10,
+            "All concurrent mTLS requests should succeed"
+        );
+
+        // Cleanup
+        worker.stop().await;
+    }
+}
+
+#[cfg(test)]
+mod certificate_tests {
+    use super::*;
+
+    /// Test that certificate generation works correctly
+    #[test]
+    fn test_certificate_generation() {
+        let certs = TestCertificates::generate().expect("Failed to generate certificates");
+
+        // Verify all files exist and are not empty
+        assert!(certs.ca_cert_path.exists());
+        assert!(certs.ca_key_path.exists());
+        assert!(certs.server_cert_path.exists());
+        assert!(certs.server_key_path.exists());
+        assert!(certs.client_cert_path.exists());
+        assert!(certs.client_key_path.exists());
+
+        // Verify PEM format
+        let ca_cert_content = std::fs::read_to_string(&certs.ca_cert_path).unwrap();
+        assert!(ca_cert_content.contains("-----BEGIN CERTIFICATE-----"));
+        assert!(ca_cert_content.contains("-----END CERTIFICATE-----"));
+
+        let server_key_content = std::fs::read_to_string(&certs.server_key_path).unwrap();
+        assert!(server_key_content.contains("-----BEGIN PRIVATE KEY-----"));
+        assert!(server_key_content.contains("-----END PRIVATE KEY-----"));
+    }
+
+    /// Test that generated certificates can be parsed by rustls
+    #[test]
+    fn test_certificate_parsing() {
+        ensure_crypto_provider();
+
+        let test_certs = TestCertificates::generate().expect("Failed to generate certificates");
+
+        // Verify CA certificate can be parsed by rustls
+        let ca_file = std::fs::File::open(&test_certs.ca_cert_path).unwrap();
+        let mut ca_reader = BufReader::new(ca_file);
+        let ca_certs: Vec<_> = certs(&mut ca_reader).filter_map(|r| r.ok()).collect();
+        assert!(!ca_certs.is_empty(), "CA certificate should be parseable");
+
+        // Verify client certificate can be parsed
+        let cert_file = std::fs::File::open(&test_certs.client_cert_path).unwrap();
+        let mut cert_reader = BufReader::new(cert_file);
+        let client_certs: Vec<_> = certs(&mut cert_reader).filter_map(|r| r.ok()).collect();
+        assert!(
+            !client_certs.is_empty(),
+            "Client certificate should be parseable"
+        );
+
+        // Verify client key can be parsed
+        let key_file = std::fs::File::open(&test_certs.client_key_path).unwrap();
+        let mut key_reader = BufReader::new(key_file);
+        let keys: Vec<_> = pkcs8_private_keys(&mut key_reader)
+            .filter_map(|r| r.ok())
+            .collect();
+        assert!(!keys.is_empty(), "Client key should be parseable");
+
+        // Verify we can build a complete client config
+        let mut root_store = RootCertStore::empty();
+        for cert in ca_certs {
+            root_store
+                .add(cert)
+                .expect("Should be able to add CA cert to store");
+        }
+
+        let key = keys.into_iter().next().unwrap();
+        let config_result = ClientConfig::builder()
+            .with_root_certificates(root_store)
+            .with_client_auth_cert(client_certs, key.into());
+
+        assert!(
+            config_result.is_ok(),
+            "Should be able to create client config with certs"
+        );
+    }
+}

sgl-model-gateway/tests/security_tests.rs

@@ -0,0 +1,8 @@
+//! Security integration tests
+
+#[path = "common/mod.rs"]
+pub mod common;
+
+mod security;
+
+pub use security::*;

sgl-model-gateway/tests/service_discovery_test.rs

@@ -0,0 +1,261 @@
+//! Service discovery integration tests
+//!
+//! Tests for service discovery shim functionality for dynamic worker registration.
+
+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 service_discovery_tests {
+    use super::*;
+
+    /// Test service discovery endpoint responds correctly
+    #[tokio::test]
+    async fn test_service_discovery_endpoint() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(4000)
+            .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)
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![MockWorkerConfig {
+                port: 20000,
+                worker_type: WorkerType::Regular,
+                health_status: HealthStatus::Healthy,
+                response_delay_ms: 0,
+                fail_rate: 0.0,
+            }],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Check if service discovery endpoint exists
+        let req = Request::builder()
+            .method("GET")
+            .uri("/v1/workers")
+            .body(Body::empty())
+            .unwrap();
+
+        let resp = app.clone().oneshot(req).await.unwrap();
+        // Endpoint might return OK with worker list or 404 if not implemented
+        assert!(
+            resp.status() == StatusCode::OK || resp.status() == StatusCode::NOT_FOUND,
+            "Workers endpoint should respond, got {}",
+            resp.status()
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test worker registration via discovery shim
+    #[tokio::test]
+    async fn test_worker_registration() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(4001)
+            .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)
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![MockWorkerConfig {
+                port: 20001,
+                worker_type: WorkerType::Regular,
+                health_status: HealthStatus::Healthy,
+                response_delay_ms: 0,
+                fail_rate: 0.0,
+            }],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Register a new worker via discovery endpoint
+        let register_payload = json!({
+            "url": "http://127.0.0.1:20002",
+            "weight": 1.0
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/register_worker")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(
+                serde_json::to_string(&register_payload).unwrap(),
+            ))
+            .unwrap();
+
+        let resp = app.clone().oneshot(req).await.unwrap();
+        // Registration might succeed or endpoint might not exist
+        assert!(
+            resp.status() == StatusCode::OK
+                || resp.status() == StatusCode::ACCEPTED
+                || resp.status() == StatusCode::NOT_FOUND,
+            "Registration should respond appropriately, got {}",
+            resp.status()
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test worker deregistration via discovery shim
+    #[tokio::test]
+    async fn test_worker_deregistration() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(4002)
+            .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)
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                MockWorkerConfig {
+                    port: 20003,
+                    worker_type: WorkerType::Regular,
+                    health_status: HealthStatus::Healthy,
+                    response_delay_ms: 0,
+                    fail_rate: 0.0,
+                },
+                MockWorkerConfig {
+                    port: 20004,
+                    worker_type: WorkerType::Regular,
+                    health_status: HealthStatus::Healthy,
+                    response_delay_ms: 0,
+                    fail_rate: 0.0,
+                },
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Deregister a worker via discovery endpoint
+        let deregister_payload = json!({
+            "url": "http://127.0.0.1:20003"
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/deregister_worker")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(
+                serde_json::to_string(&deregister_payload).unwrap(),
+            ))
+            .unwrap();
+
+        let resp = app.clone().oneshot(req).await.unwrap();
+        assert!(
+            resp.status() == StatusCode::OK
+                || resp.status() == StatusCode::ACCEPTED
+                || resp.status() == StatusCode::NOT_FOUND,
+            "Deregistration should respond appropriately, got {}",
+            resp.status()
+        );
+
+        // Requests should still work with remaining worker
+        let payload = json!({
+            "text": "Test after deregistration",
+            "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,
+            "Request should succeed with remaining worker"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test health status reporting for discovery
+    #[tokio::test]
+    async fn test_health_status_endpoint() {
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .round_robin_policy()
+            .host("127.0.0.1")
+            .port(4003)
+            .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)
+            .build_unchecked();
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![MockWorkerConfig {
+                port: 20005,
+                worker_type: WorkerType::Regular,
+                health_status: HealthStatus::Healthy,
+                response_delay_ms: 0,
+                fail_rate: 0.0,
+            }],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Check health endpoint
+        let req = Request::builder()
+            .method("GET")
+            .uri("/health")
+            .body(Body::empty())
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "Health endpoint should return OK when workers are healthy"
+        );
+
+        ctx.shutdown().await;
+    }
+}

sgl-model-gateway/tests/tool_parser/mod.rs

@@ -0,0 +1,17 @@
+//! Tool parser integration tests
+
+pub mod tool_parser_deepseek;
+pub mod tool_parser_edge_cases;
+pub mod tool_parser_fallback;
+pub mod tool_parser_glm47_moe;
+pub mod tool_parser_glm4_moe;
+pub mod tool_parser_json;
+pub mod tool_parser_kimik2;
+pub mod tool_parser_llama;
+pub mod tool_parser_minimax_m2;
+pub mod tool_parser_mistral;
+pub mod tool_parser_mixed_edge_cases;
+pub mod tool_parser_partial_json;
+pub mod tool_parser_pythonic;
+pub mod tool_parser_qwen;
+pub mod tool_parser_step3;

sgl-model-gateway/tests/tool_parser/tool_parser_deepseek.rs

@@ -2,8 +2,7 @@
 
 use smg::tool_parser::{DeepSeekParser, ToolParser};
 
-mod common;
-use common::create_test_tools;
+use crate::common::create_test_tools;
 
 #[tokio::test]
 async fn test_deepseek_complete_parsing() {

sgl-model-gateway/tests/tool_parser/tool_parser_edge_cases.rs

@@ -4,8 +4,7 @@
 
 use smg::tool_parser::{JsonParser, MistralParser, PythonicParser, QwenParser, ToolParser};
 
-mod common;
-use common::create_test_tools;
+use crate::common::create_test_tools;
 
 #[tokio::test]
 async fn test_empty_input() {

sgl-model-gateway/tests/tool_parser/tool_parser_glm47_moe.rs

@@ -2,8 +2,7 @@
 
 use smg::tool_parser::{Glm4MoeParser, ToolParser};
 
-mod common;
-use common::create_test_tools;
+use crate::common::create_test_tools;
 
 #[tokio::test]
 async fn test_glm47_complete_parsing() {

sgl-model-gateway/tests/tool_parser/tool_parser_glm4_moe.rs

@@ -2,8 +2,7 @@
 
 use smg::tool_parser::{Glm4MoeParser, ToolParser};
 
-mod common;
-use common::create_test_tools;
+use crate::common::create_test_tools;
 
 #[tokio::test]
 async fn test_glm4_complete_parsing() {

sgl-model-gateway/tests/tool_parser/tool_parser_json.rs

@@ -5,8 +5,7 @@
 use serde_json::json;
 use smg::tool_parser::{JsonParser, ToolParser};
 
-mod common;
-use common::{create_test_tools, streaming_helpers::*};
+use crate::common::{create_test_tools, streaming_helpers::*};
 
 #[tokio::test]
 async fn test_simple_json_tool_call() {

sgl-model-gateway/tests/tool_parser/tool_parser_kimik2.rs

@@ -2,8 +2,7 @@
 
 use smg::tool_parser::{KimiK2Parser, ToolParser};
 
-mod common;
-use common::create_test_tools;
+use crate::common::create_test_tools;
 
 #[tokio::test]
 async fn test_kimik2_complete_parsing() {

sgl-model-gateway/tests/tool_parser/tool_parser_llama.rs

@@ -4,8 +4,7 @@
 
 use smg::tool_parser::{LlamaParser, ToolParser};
 
-mod common;
-use common::{create_test_tools, streaming_helpers::*};
+use crate::common::{create_test_tools, streaming_helpers::*};
 
 #[tokio::test]
 async fn test_llama_python_tag_format() {

sgl-model-gateway/tests/tool_parser/tool_parser_minimax_m2.rs

@@ -2,8 +2,7 @@
 
 use smg::tool_parser::{MinimaxM2Parser, ToolParser};
 
-mod common;
-use common::create_test_tools;
+use crate::common::create_test_tools;
 
 #[tokio::test]
 async fn test_minimax_complete_parsing() {

sgl-model-gateway/tests/tool_parser/tool_parser_mixed_edge_cases.rs

@@ -7,8 +7,7 @@ use smg::tool_parser::{
     JsonParser, LlamaParser, MistralParser, PythonicParser, QwenParser, ToolParser,
 };
 
-mod common;
-use common::create_test_tools;
+use crate::common::create_test_tools;
 
 #[tokio::test]
 async fn test_mixed_formats_in_text() {

sgl-model-gateway/tests/tool_parser/tool_parser_pythonic.rs

@@ -5,8 +5,7 @@
 use serde_json::json;
 use smg::tool_parser::{PythonicParser, ToolParser};
 
-mod common;
-use common::create_test_tools;
+use crate::common::create_test_tools;
 
 #[tokio::test]
 async fn test_pythonic_single_function() {

sgl-model-gateway/tests/tool_parser/tool_parser_qwen.rs

@@ -5,8 +5,7 @@
 use serde_json::json;
 use smg::tool_parser::{QwenParser, ToolParser};
 
-mod common;
-use common::{create_test_tools, streaming_helpers::*};
+use crate::common::{create_test_tools, streaming_helpers::*};
 
 #[tokio::test]
 async fn test_qwen_single_tool() {

sgl-model-gateway/tests/tool_parser/tool_parser_step3.rs

@@ -2,8 +2,7 @@
 
 use smg::tool_parser::{Step3Parser, ToolParser};
 
-mod common;
-use common::create_test_tools;
+use crate::common::create_test_tools;
 
 #[tokio::test]
 async fn test_step3_complete_parsing() {

sgl-model-gateway/tests/tool_parser_tests.rs

@@ -0,0 +1,8 @@
+//! Tool parser integration tests
+
+#[path = "common/mod.rs"]
+pub mod common;
+
+mod tool_parser;
+
+pub use tool_parser::*;

sgl-model-gateway/tests/worker_management_test.rs

@@ -0,0 +1,140 @@
+//! Worker management integration tests
+//!
+//! Tests for dynamic worker add/remove operations via management API.
+//! The actual worker management API uses:
+//! - POST /workers - create a worker
+//! - GET /workers - list workers
+//! - DELETE /workers/{worker_id} - remove a worker
+
+mod common;
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
+use serde_json::json;
+use tower::ServiceExt;
+
+#[cfg(test)]
+mod worker_management_tests {
+    use super::*;
+
+    /// Test listing workers via API
+    #[tokio::test]
+    async fn test_list_workers() {
+        let config = TestRouterConfig::round_robin(3900);
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::healthy(19900),
+                TestWorkerConfig::healthy(19901),
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // List workers via GET /workers
+        let req = Request::builder()
+            .method("GET")
+            .uri("/workers")
+            .body(Body::empty())
+            .unwrap();
+
+        let resp = app.clone().oneshot(req).await.unwrap();
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "GET /workers should return OK"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test that routing continues to work with multiple workers
+    #[tokio::test]
+    async fn test_routing_with_multiple_workers() {
+        let config = TestRouterConfig::round_robin(3901);
+
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::healthy(19902),
+                TestWorkerConfig::healthy(19903),
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+        let mut success_count = 0;
+
+        // Verify routing distributes across workers
+        for i in 0..20 {
+            let payload = json!({
+                "text": format!("Test request {}", i),
+                "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.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        assert_eq!(
+            success_count, 20,
+            "All requests should succeed with multiple workers"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test that requests continue to work during worker operations
+    #[tokio::test]
+    async fn test_requests_during_worker_changes() {
+        let config = TestRouterConfig::round_robin(3902);
+
+        let ctx =
+            AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(19904)]).await;
+
+        let app = ctx.create_app().await;
+
+        // Send requests and verify they succeed
+        let mut success_count = 0;
+        for i in 0..10 {
+            let payload = json!({
+                "text": format!("Request during changes {}", i),
+                "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.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        assert_eq!(
+            success_count, 10,
+            "All requests should succeed during normal operation"
+        );
+
+        ctx.shutdown().await;
+    }
+}