[model-gateway] reorganize integration tests into logical subdirectories (#16451)

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

@@ -0,0 +1,308 @@
+//! Header forwarding integration tests
+//!
+//! Tests for header propagation through the router to workers.
+
+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,
+};
+
+#[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/routing/mod.rs

@@ -1,9 +1,14 @@
 //! Routing integration tests
 
 pub mod cache_aware_backward_compat_test;
+pub mod header_forwarding_test;
 pub mod load_balancing_test;
 pub mod manual_routing_test;
+pub mod payload_size_test;
 pub mod pd_routing_test;
+pub mod policy_registry_integration;
 pub mod power_of_two_test;
+pub mod service_discovery_test;
 pub mod test_openai_routing;
 pub mod test_pd_routing;
+pub mod worker_management_test;

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

@@ -0,0 +1,297 @@
+//! Payload size integration tests
+//!
+//! Tests for request payload size limits and handling.
+
+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,
+};
+
+#[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/routing/policy_registry_integration.rs

@@ -0,0 +1,170 @@
+//! Integration tests for PolicyRegistry with RouterManager
+
+use std::{collections::HashMap, sync::Arc};
+
+use smg::{
+    config::PolicyConfig, core::WorkerRegistry, policies::PolicyRegistry,
+    protocols::worker_spec::WorkerConfigRequest, routers::router_manager::RouterManager,
+};
+
+#[tokio::test]
+async fn test_policy_registry_with_router_manager() {
+    // Create HTTP client
+    let _client = reqwest::Client::new();
+
+    // Create shared registries
+    let worker_registry = Arc::new(WorkerRegistry::new());
+    let policy_registry = Arc::new(PolicyRegistry::new(PolicyConfig::RoundRobin));
+
+    // Create RouterManager with shared registries
+    let _router_manager = RouterManager::new(worker_registry.clone());
+
+    // Add first worker for llama-3 with cache_aware policy hint
+    let mut labels1 = HashMap::new();
+    labels1.insert("policy".to_string(), "cache_aware".to_string());
+
+    let _worker1_config = WorkerConfigRequest {
+        url: "http://worker1:8000".to_string(),
+        model_id: Some("llama-3".to_string()),
+        api_key: Some("test_api_key".to_string()),
+        worker_type: None,
+        priority: None,
+        cost: None,
+        labels: labels1,
+        bootstrap_port: None,
+        tokenizer_path: None,
+        reasoning_parser: None,
+        tool_parser: None,
+        chat_template: None,
+        runtime: None,
+        health_check_timeout_secs: 30,
+        health_check_interval_secs: 60,
+        health_success_threshold: 2,
+        health_failure_threshold: 3,
+        max_connection_attempts: 20,
+        dp_aware: false,
+    };
+
+    // This would normally connect to a real worker, but for testing we'll just verify the structure
+    // In a real test, we'd need to mock the worker or use a test server
+
+    let _llama_policy = policy_registry.get_policy("llama-3");
+    // After first worker is added, llama-3 should have a policy
+
+    // Add second worker for llama-3 with different policy hint (should be ignored)
+    let mut labels2 = HashMap::new();
+    labels2.insert("policy".to_string(), "random".to_string());
+
+    let _worker2_config = WorkerConfigRequest {
+        url: "http://worker2:8000".to_string(),
+        model_id: Some("llama-3".to_string()),
+        api_key: Some("test_api_key".to_string()),
+        worker_type: None,
+        priority: None,
+        cost: None,
+        labels: labels2,
+        bootstrap_port: None,
+        tokenizer_path: None,
+        reasoning_parser: None,
+        tool_parser: None,
+        chat_template: None,
+        runtime: None,
+        health_check_timeout_secs: 30,
+        health_check_interval_secs: 60,
+        health_success_threshold: 2,
+        health_failure_threshold: 3,
+        max_connection_attempts: 20,
+        dp_aware: false,
+    };
+
+    // The second worker should use the same policy as the first (cache_aware)
+
+    // Add worker for different model (gpt-4) with random policy
+    let mut labels3 = HashMap::new();
+    labels3.insert("policy".to_string(), "random".to_string());
+
+    let _worker3_config = WorkerConfigRequest {
+        url: "http://worker3:8000".to_string(),
+        model_id: Some("gpt-4".to_string()),
+        api_key: Some("test_api_key".to_string()),
+        worker_type: None,
+        priority: None,
+        cost: None,
+        labels: labels3,
+        bootstrap_port: None,
+        tokenizer_path: None,
+        reasoning_parser: None,
+        tool_parser: None,
+        runtime: None,
+        chat_template: None,
+        health_check_timeout_secs: 30,
+        health_check_interval_secs: 60,
+        health_success_threshold: 2,
+        health_failure_threshold: 3,
+        max_connection_attempts: 20,
+        dp_aware: false,
+    };
+
+    let _gpt_policy = policy_registry.get_policy("gpt-4");
+
+    // When we remove both llama-3 workers, the policy should be cleaned up
+
+    println!("PolicyRegistry integration test structure created");
+    println!("Note: This test requires mocking or test servers to fully execute");
+}
+
+#[test]
+fn test_policy_registry_cleanup() {
+    use smg::{config::PolicyConfig, policies::PolicyRegistry};
+
+    let registry = PolicyRegistry::new(PolicyConfig::RoundRobin);
+
+    // Add workers for a model
+    let policy1 = registry.on_worker_added("model-1", Some("cache_aware"));
+    assert_eq!(policy1.name(), "cache_aware");
+
+    // Second worker uses existing policy
+    let policy2 = registry.on_worker_added("model-1", Some("random"));
+    assert_eq!(policy2.name(), "cache_aware"); // Should still be cache_aware
+
+    assert!(registry.get_policy("model-1").is_some());
+
+    // Remove first worker - policy should remain
+    registry.on_worker_removed("model-1");
+    assert!(registry.get_policy("model-1").is_some());
+
+    // Remove second worker - policy should be cleaned up
+    registry.on_worker_removed("model-1");
+    assert!(registry.get_policy("model-1").is_none());
+
+    println!("✓ PolicyRegistry cleanup test passed");
+}
+
+#[test]
+fn test_policy_registry_multiple_models() {
+    use smg::{config::PolicyConfig, policies::PolicyRegistry};
+
+    let registry = PolicyRegistry::new(PolicyConfig::RoundRobin);
+
+    // Add workers for different models with different policies
+    let llama_policy = registry.on_worker_added("llama-3", Some("cache_aware"));
+    let gpt_policy = registry.on_worker_added("gpt-4", Some("random"));
+    let mistral_policy = registry.on_worker_added("mistral", None); // Uses default
+
+    assert_eq!(llama_policy.name(), "cache_aware");
+    assert_eq!(gpt_policy.name(), "random");
+    assert_eq!(mistral_policy.name(), "round_robin"); // Default
+
+    assert!(registry.get_policy("llama-3").is_some());
+    assert!(registry.get_policy("gpt-4").is_some());
+    assert!(registry.get_policy("mistral").is_some());
+
+    // Get all mappings
+    let mappings = registry.get_all_mappings();
+    assert_eq!(mappings.len(), 3);
+    assert_eq!(mappings.get("llama-3").unwrap(), "cache_aware");
+    assert_eq!(mappings.get("gpt-4").unwrap(), "random");
+    assert_eq!(mappings.get("mistral").unwrap(), "round_robin");
+
+    println!("✓ PolicyRegistry multiple models test passed");
+}

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

@@ -0,0 +1,260 @@
+//! Service discovery integration tests
+//!
+//! Tests for service discovery shim functionality for dynamic worker registration.
+
+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,
+};
+
+#[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/routing/worker_management_test.rs

@@ -0,0 +1,139 @@
+//! 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
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use serde_json::json;
+use tower::ServiceExt;
+
+use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
+
+#[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;
+    }
+}