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

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