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

2026-01-04 21:49:20 -08:00
parent 4ea6a11c83
commit e53160bb31
24 changed files with 99 additions and 65 deletions
@@ -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;
+    }
+}