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

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;
+    }
+}