Support min num routing keys in key-based load balancing policy (#16564)

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

@@ -299,10 +299,12 @@ async fn generate_handler(
     Json(payload): Json<serde_json::Value>,
 ) -> Response {
     let config = config.read().await;
+    let worker_id = format!("worker-{}", config.port);
 
     if should_fail(&config).await {
         return (
             StatusCode::INTERNAL_SERVER_ERROR,
+            [("x-worker-id", worker_id)],
             Json(json!({
                 "error": "Random failure for testing"
             })),
@@ -373,28 +375,33 @@ async fn generate_handler(
 
         let stream = stream::iter(events);
 
-        Sse::new(stream)
-            .keep_alive(KeepAlive::default())
+        (
+            [("x-worker-id", worker_id)],
+            Sse::new(stream).keep_alive(KeepAlive::default()),
+        )
             .into_response()
     } else {
-        Json(json!({
-            "text": "This is a mock response.",
-            "meta_info": {
-                "prompt_tokens": 10,
-                "completion_tokens": 5,
-                "completion_tokens_wo_jump_forward": 5,
-                "input_token_logprobs": null,
-                "output_token_logprobs": null,
-                "first_token_latency": config.response_delay_ms as f64 / 1000.0,
-                "time_to_first_token": config.response_delay_ms as f64 / 1000.0,
-                "time_per_output_token": 0.01,
-                "finish_reason": {
-                    "type": "stop",
-                    "reason": "length"
+        (
+            [("x-worker-id", worker_id)],
+            Json(json!({
+                "text": "This is a mock response.",
+                "meta_info": {
+                    "prompt_tokens": 10,
+                    "completion_tokens": 5,
+                    "completion_tokens_wo_jump_forward": 5,
+                    "input_token_logprobs": null,
+                    "output_token_logprobs": null,
+                    "first_token_latency": config.response_delay_ms as f64 / 1000.0,
+                    "time_to_first_token": config.response_delay_ms as f64 / 1000.0,
+                    "time_per_output_token": 0.01,
+                    "finish_reason": {
+                        "type": "stop",
+                        "reason": "length"
+                    }
                 }
-            }
-        }))
-        .into_response()
+            })),
+        )
+            .into_response()
     }
 }
 

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

@@ -3,7 +3,9 @@
 //! Provides pre-configured RouterConfig and MockWorkerConfig builders
 //! for common test scenarios.
 
-use smg::config::{CircuitBreakerConfig, PolicyConfig, RetryConfig, RouterConfig};
+use smg::config::{
+    CircuitBreakerConfig, ManualAssignmentMode, PolicyConfig, RetryConfig, RouterConfig,
+};
 
 use super::mock_worker::{HealthStatus, MockWorkerConfig, WorkerType};
 
@@ -94,12 +96,22 @@ impl TestRouterConfig {
 
     /// Create a manual routing config (for sticky routing tests)
     pub fn manual(port: u16) -> RouterConfig {
+        Self::manual_with_mode(port, ManualAssignmentMode::Random)
+    }
+
+    /// Create a manual routing config with min_group assignment mode
+    pub fn manual_min_group(port: u16) -> RouterConfig {
+        Self::manual_with_mode(port, ManualAssignmentMode::MinGroup)
+    }
+
+    /// Create a manual routing config with specified assignment mode
+    pub fn manual_with_mode(port: u16, assignment_mode: ManualAssignmentMode) -> RouterConfig {
         RouterConfig::builder()
             .regular_mode(vec![])
             .policy(PolicyConfig::Manual {
                 eviction_interval_secs: 60,
                 max_idle_secs: 3600,
-                assignment_mode: Default::default(),
+                assignment_mode,
             })
             .host(defaults::HOST)
             .port(port)
@@ -262,6 +274,13 @@ impl TestWorkerConfig {
         }
     }
 
+    /// Create multiple slow workers with sequential ports
+    pub fn slow_workers(start_port: u16, count: u16, delay_ms: u64) -> Vec<MockWorkerConfig> {
+        (0..count)
+            .map(|i| Self::slow(start_port + i, delay_ms))
+            .collect()
+    }
+
     /// Create a flaky worker config (for retry/fault tolerance tests)
     pub fn flaky(port: u16, fail_rate: f32) -> MockWorkerConfig {
         MockWorkerConfig {