api_endpoints_test.rs 50.7 KB
Newer Older
1
2
mod common;

3
4
5
6
7
use axum::{
    body::Body,
    extract::Request,
    http::{header::CONTENT_TYPE, StatusCode},
};
8
9
10
use common::mock_worker::{HealthStatus, MockWorker, MockWorkerConfig, WorkerType};
use reqwest::Client;
use serde_json::json;
11
12
13
use sglang_router_rs::config::{
    CircuitBreakerConfig, PolicyConfig, RetryConfig, RouterConfig, RoutingMode,
};
14
15
16
use sglang_router_rs::routers::{RouterFactory, RouterTrait};
use std::sync::Arc;
use tower::ServiceExt;
17
18
19
20

/// Test context that manages mock workers
struct TestContext {
    workers: Vec<MockWorker>,
21
22
23
    router: Arc<dyn RouterTrait>,
    client: Client,
    config: RouterConfig,
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
}

impl TestContext {
    async fn new(worker_configs: Vec<MockWorkerConfig>) -> Self {
        // Create default router config
        let config = RouterConfig {
            mode: RoutingMode::Regular {
                worker_urls: vec![],
            },
            policy: PolicyConfig::Random,
            host: "127.0.0.1".to_string(),
            port: 3002,
            max_payload_size: 256 * 1024 * 1024,
            request_timeout_secs: 600,
            worker_startup_timeout_secs: 1,
            worker_startup_check_interval_secs: 1,
40
            discovery: None,
41
42
            dp_aware: false,
            api_key: None,
43
44
45
            metrics: None,
            log_dir: None,
            log_level: None,
46
            request_id_headers: None,
47
            max_concurrent_requests: 64,
48
49
50
            queue_size: 0,
            queue_timeout_secs: 60,
            rate_limit_tokens_per_second: None,
51
            cors_allowed_origins: vec![],
52
            retry: RetryConfig::default(),
53
            circuit_breaker: CircuitBreakerConfig::default(),
54
55
            disable_retries: false,
            disable_circuit_breaker: false,
56
            health_check: sglang_router_rs::config::HealthCheckConfig::default(),
57
            enable_igw: false,
58
59
60
61
62
        };

        Self::new_with_config(config, worker_configs).await
    }

63
64
65
66
    async fn new_with_config(
        mut config: RouterConfig,
        worker_configs: Vec<MockWorkerConfig>,
    ) -> Self {
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
        let mut workers = Vec::new();
        let mut worker_urls = Vec::new();

        // Start mock workers if any
        for worker_config in worker_configs {
            let mut worker = MockWorker::new(worker_config);
            let url = worker.start().await.unwrap();
            worker_urls.push(url);
            workers.push(worker);
        }

        if !workers.is_empty() {
            tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;
        }

82
83
84
85
86
87
88
89
90
91
        // Update config with worker URLs if not already set
        if let RoutingMode::Regular {
            worker_urls: ref mut urls,
        } = config.mode
        {
            if urls.is_empty() {
                *urls = worker_urls.clone();
            }
        }

92
93
94
95
96
        let client = Client::builder()
            .timeout(std::time::Duration::from_secs(config.request_timeout_secs))
            .build()
            .unwrap();

97
98
        // Create app context
        let app_context = common::create_test_context(config.clone());
99

100
101
        // Create router
        let router = RouterFactory::create_router(&app_context).await.unwrap();
102
        let router = Arc::from(router);
103

104
105
        // Wait for router to discover workers
        if !workers.is_empty() {
106
107
108
            tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
        }

109
110
111
112
113
114
        Self {
            workers,
            router,
            client,
            config,
        }
115
116
    }

117
118
119
120
121
    async fn create_app(&self) -> axum::Router {
        common::test_app::create_test_app(
            Arc::clone(&self.router),
            self.client.clone(),
            &self.config,
122
123
124
125
126
127
128
129
130
131
132
133
134
135
        )
    }

    async fn shutdown(mut self) {
        for worker in &mut self.workers {
            worker.stop().await;
        }
    }
}

#[cfg(test)]
mod health_tests {
    use super::*;

136
137
138
139
    #[tokio::test]
    async fn test_liveness_endpoint() {
        let ctx = TestContext::new(vec![]).await;
        let app = ctx.create_app().await;
140

141
142
143
144
145
        let req = Request::builder()
            .method("GET")
            .uri("/liveness")
            .body(Body::empty())
            .unwrap();
146

147
148
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
149

150
        ctx.shutdown().await;
151
152
    }

153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
    #[tokio::test]
    async fn test_readiness_with_healthy_workers() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18001,
            worker_type: WorkerType::Regular,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 0.0,
        }])
        .await;

        let app = ctx.create_app().await;

        let req = Request::builder()
            .method("GET")
            .uri("/readiness")
            .body(Body::empty())
            .unwrap();
171

172
173
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
174

175
        ctx.shutdown().await;
176
177
    }

178
179
180
    #[tokio::test]
    async fn test_readiness_with_unhealthy_workers() {
        let ctx = TestContext::new(vec![]).await;
181

182
        let app = ctx.create_app().await;
183

184
185
186
187
188
        let req = Request::builder()
            .method("GET")
            .uri("/readiness")
            .body(Body::empty())
            .unwrap();
189

190
191
192
        let resp = app.oneshot(req).await.unwrap();
        // With no workers, readiness should return SERVICE_UNAVAILABLE
        assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
193

194
        ctx.shutdown().await;
195
196
    }

197
198
199
200
201
202
203
204
205
206
207
208
    #[tokio::test]
    async fn test_health_endpoint_details() {
        let ctx = TestContext::new(vec![
            MockWorkerConfig {
                port: 18003,
                worker_type: WorkerType::Regular,
                health_status: HealthStatus::Healthy,
                response_delay_ms: 0,
                fail_rate: 0.0,
            },
            MockWorkerConfig {
                port: 18004,
209
210
211
212
                worker_type: WorkerType::Regular,
                health_status: HealthStatus::Healthy,
                response_delay_ms: 0,
                fail_rate: 0.0,
213
214
215
            },
        ])
        .await;
216

217
        let app = ctx.create_app().await;
218

219
220
221
222
223
        let req = Request::builder()
            .method("GET")
            .uri("/health")
            .body(Body::empty())
            .unwrap();
224

225
226
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
227

228
229
230
231
232
233
        // The health endpoint returns plain text, not JSON
        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        let body_str = String::from_utf8_lossy(&body);
        assert!(body_str.contains("All servers healthy"));
234

235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
        ctx.shutdown().await;
    }

    #[tokio::test]
    async fn test_health_generate_endpoint() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18005,
            worker_type: WorkerType::Regular,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 0.0,
        }])
        .await;

        let app = ctx.create_app().await;

        let req = Request::builder()
            .method("GET")
            .uri("/health_generate")
            .body(Body::empty())
            .unwrap();

        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);

        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert!(body_json.is_object());

        ctx.shutdown().await;
267
268
269
270
271
272
273
    }
}

#[cfg(test)]
mod generation_tests {
    use super::*;

274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
    #[tokio::test]
    async fn test_generate_success() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18101,
            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": "Hello, world!",
            "stream": false
        });
291

292
293
294
295
296
297
        let req = Request::builder()
            .method("POST")
            .uri("/generate")
            .header(CONTENT_TYPE, "application/json")
            .body(Body::from(serde_json::to_string(&payload).unwrap()))
            .unwrap();
298

299
300
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
301

302
303
304
305
306
307
308
309
310
311
312
313
        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert!(body_json.get("text").is_some());
        assert!(body_json.get("meta_info").is_some());
        let meta_info = &body_json["meta_info"];
        assert!(meta_info.get("finish_reason").is_some());
        assert_eq!(meta_info["finish_reason"]["type"], "stop");

        ctx.shutdown().await;
    }
314

315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
    #[tokio::test]
    async fn test_generate_streaming() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18102,
            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": "Stream test",
            "stream": true
331
332
        });

333
334
335
336
337
338
        let req = Request::builder()
            .method("POST")
            .uri("/generate")
            .header(CONTENT_TYPE, "application/json")
            .body(Body::from(serde_json::to_string(&payload).unwrap()))
            .unwrap();
339

340
341
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
342

343
344
345
346
        // For streaming responses, the router might use chunked encoding or other streaming mechanisms
        // The exact content-type can vary based on the router implementation
        // Just verify we got a successful response
        // Note: In a real implementation, we'd check for text/event-stream or appropriate streaming headers
347

348
349
        ctx.shutdown().await;
    }
350

351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
    #[tokio::test]
    async fn test_generate_with_worker_failure() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18103,
            worker_type: WorkerType::Regular,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 1.0, // Always fail
        }])
        .await;

        let app = ctx.create_app().await;

        let payload = json!({
            "text": "This should fail",
            "stream": false
        });
368

369
370
371
372
373
374
        let req = Request::builder()
            .method("POST")
            .uri("/generate")
            .header(CONTENT_TYPE, "application/json")
            .body(Body::from(serde_json::to_string(&payload).unwrap()))
            .unwrap();
375

376
377
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
378

379
380
        ctx.shutdown().await;
    }
381

382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
    #[tokio::test]
    async fn test_v1_chat_completions_success() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18104,
            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
        });
402

403
404
405
406
407
408
        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();
409

410
411
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
412

413
414
415
416
417
        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert!(body_json.get("choices").is_some());
418

419
        ctx.shutdown().await;
420
    }
421
}
422

423
424
425
#[cfg(test)]
mod model_info_tests {
    use super::*;
426

427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
    #[tokio::test]
    async fn test_get_server_info() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18201,
            worker_type: WorkerType::Regular,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 0.0,
        }])
        .await;

        let app = ctx.create_app().await;

        let req = Request::builder()
            .method("GET")
            .uri("/get_server_info")
            .body(Body::empty())
            .unwrap();
445

446
447
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
448

449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert!(body_json.is_object());
        // Check for actual sglang server fields
        assert!(body_json.get("version").is_some());
        assert!(body_json.get("model_path").is_some());
        assert!(body_json.get("tokenizer_path").is_some());
        assert!(body_json.get("port").is_some());
        assert!(body_json.get("max_num_batched_tokens").is_some());
        assert!(body_json.get("schedule_policy").is_some());

        ctx.shutdown().await;
    }
464

465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
    #[tokio::test]
    async fn test_get_model_info() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18202,
            worker_type: WorkerType::Regular,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 0.0,
        }])
        .await;

        let app = ctx.create_app().await;

        let req = Request::builder()
            .method("GET")
            .uri("/get_model_info")
            .body(Body::empty())
            .unwrap();
483

484
485
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
486

487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert!(body_json.is_object());
        // Check for actual sglang model info fields
        assert_eq!(
            body_json.get("model_path").and_then(|v| v.as_str()),
            Some("mock-model-path")
        );
        assert_eq!(
            body_json.get("tokenizer_path").and_then(|v| v.as_str()),
            Some("mock-tokenizer-path")
        );
        assert_eq!(
            body_json.get("is_generation").and_then(|v| v.as_bool()),
            Some(true)
        );
        assert!(body_json.get("preferred_sampling_params").is_some());

        ctx.shutdown().await;
508
509
    }

510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
    #[tokio::test]
    async fn test_v1_models() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18203,
            worker_type: WorkerType::Regular,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 0.0,
        }])
        .await;

        let app = ctx.create_app().await;

        let req = Request::builder()
            .method("GET")
            .uri("/v1/models")
            .body(Body::empty())
            .unwrap();
528

529
530
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
531

532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert!(body_json.get("object").is_some());
        assert_eq!(
            body_json.get("object").and_then(|v| v.as_str()),
            Some("list")
        );

        let data = body_json.get("data").and_then(|v| v.as_array());
        assert!(data.is_some());

        let models = data.unwrap();
        assert!(!models.is_empty());

        let first_model = &models[0];
        assert_eq!(
            first_model.get("id").and_then(|v| v.as_str()),
            Some("mock-model")
        );
        assert_eq!(
            first_model.get("object").and_then(|v| v.as_str()),
            Some("model")
        );
        assert!(first_model.get("created").is_some());
        assert_eq!(
            first_model.get("owned_by").and_then(|v| v.as_str()),
            Some("organization-owner")
        );

        ctx.shutdown().await;
    }
565

566
567
568
569
    #[tokio::test]
    async fn test_model_info_with_no_workers() {
        let ctx = TestContext::new(vec![]).await;
        let app = ctx.create_app().await;
570

571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
        // Test server info with no workers
        let req = Request::builder()
            .method("GET")
            .uri("/get_server_info")
            .body(Body::empty())
            .unwrap();
        let resp = app.clone().oneshot(req).await.unwrap();
        // Router may return various error codes when no workers
        assert!(
            resp.status() == StatusCode::OK
                || resp.status() == StatusCode::SERVICE_UNAVAILABLE
                || resp.status() == StatusCode::NOT_FOUND
                || resp.status() == StatusCode::INTERNAL_SERVER_ERROR,
            "Unexpected status code: {:?}",
            resp.status()
        );

        // Test model info with no workers
        let req = Request::builder()
            .method("GET")
            .uri("/get_model_info")
            .body(Body::empty())
            .unwrap();
        let resp = app.clone().oneshot(req).await.unwrap();
        // Router may return various error codes when no workers
        assert!(
            resp.status() == StatusCode::OK
                || resp.status() == StatusCode::SERVICE_UNAVAILABLE
                || resp.status() == StatusCode::NOT_FOUND
                || resp.status() == StatusCode::INTERNAL_SERVER_ERROR,
            "Unexpected status code: {:?}",
            resp.status()
        );

        // Test v1/models with no workers
        let req = Request::builder()
            .method("GET")
            .uri("/v1/models")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        // Router may return various error codes when no workers
        assert!(
            resp.status() == StatusCode::OK
                || resp.status() == StatusCode::SERVICE_UNAVAILABLE
                || resp.status() == StatusCode::NOT_FOUND
                || resp.status() == StatusCode::INTERNAL_SERVER_ERROR,
            "Unexpected status code: {:?}",
            resp.status()
        );

        ctx.shutdown().await;
623
624
    }

625
626
627
628
629
630
631
632
633
634
635
636
    #[tokio::test]
    async fn test_model_info_with_multiple_workers() {
        let ctx = TestContext::new(vec![
            MockWorkerConfig {
                port: 18204,
                worker_type: WorkerType::Regular,
                health_status: HealthStatus::Healthy,
                response_delay_ms: 0,
                fail_rate: 0.0,
            },
            MockWorkerConfig {
                port: 18205,
637
638
639
640
                worker_type: WorkerType::Regular,
                health_status: HealthStatus::Healthy,
                response_delay_ms: 0,
                fail_rate: 0.0,
641
642
643
            },
        ])
        .await;
644

645
        let app = ctx.create_app().await;
646

647
648
649
650
        // Test that model info is consistent across workers
        for _ in 0..5 {
            let req = Request::builder()
                .method("GET")
651
                .uri("/get_model_info")
652
653
                .body(Body::empty())
                .unwrap();
654

655
            let resp = app.clone().oneshot(req).await.unwrap();
656
657
            assert_eq!(resp.status(), StatusCode::OK);

658
659
660
661
            let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
                .await
                .unwrap();
            let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
662
            assert_eq!(
663
                body_json.get("model_path").and_then(|v| v.as_str()),
664
665
                Some("mock-model-path")
            );
666
        }
667

668
        ctx.shutdown().await;
669
670
    }

671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
    #[tokio::test]
    async fn test_model_info_with_unhealthy_worker() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18206,
            worker_type: WorkerType::Regular,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 1.0, // Always fail
        }])
        .await;

        let app = ctx.create_app().await;

        let req = Request::builder()
            .method("GET")
            .uri("/get_model_info")
            .body(Body::empty())
            .unwrap();
689

690
691
692
693
694
695
696
697
        let resp = app.oneshot(req).await.unwrap();
        // Worker with fail_rate: 1.0 should always return an error status
        assert!(
            resp.status() == StatusCode::INTERNAL_SERVER_ERROR
                || resp.status() == StatusCode::SERVICE_UNAVAILABLE,
            "Expected error status for always-failing worker, got: {:?}",
            resp.status()
        );
698

699
700
701
        ctx.shutdown().await;
    }
}
702

703
704
705
#[cfg(test)]
mod worker_management_tests {
    use super::*;
706

707
708
709
710
711
712
713
714
715
716
717
718
    #[tokio::test]
    async fn test_add_new_worker() {
        let ctx = TestContext::new(vec![]).await;
        let app = ctx.create_app().await;

        // Start a mock worker
        let mut worker = MockWorker::new(MockWorkerConfig {
            port: 18301,
            worker_type: WorkerType::Regular,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 0.0,
719
        });
720
        let url = worker.start().await.unwrap();
721

722
723
724
        // Add the worker
        let req = Request::builder()
            .method("POST")
725
            .uri(format!("/add_worker?url={}", url))
726
727
            .body(Body::empty())
            .unwrap();
728

729
730
        let resp = app.clone().oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
731

732
733
734
735
736
737
        // List workers to verify
        let req = Request::builder()
            .method("GET")
            .uri("/list_workers")
            .body(Body::empty())
            .unwrap();
738

739
740
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
741

742
743
744
745
746
747
        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        let workers = body_json["urls"].as_array().unwrap();
        assert!(workers.iter().any(|w| w.as_str().unwrap() == url));
748

749
750
        worker.stop().await;
        ctx.shutdown().await;
751
752
    }

753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
    #[tokio::test]
    async fn test_remove_existing_worker() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18302,
            worker_type: WorkerType::Regular,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 0.0,
        }])
        .await;

        let app = ctx.create_app().await;

        // Get the worker URL
        let req = Request::builder()
            .method("GET")
            .uri("/list_workers")
            .body(Body::empty())
            .unwrap();
        let resp = app.clone().oneshot(req).await.unwrap();
        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        let workers = body_json["urls"].as_array().unwrap();
        let worker_url = workers[0].as_str().unwrap();

        // Remove the worker
        let req = Request::builder()
            .method("POST")
783
            .uri(format!("/remove_worker?url={}", worker_url))
784
785
            .body(Body::empty())
            .unwrap();
786

787
788
        let resp = app.clone().oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
789

790
791
792
793
794
795
796
797
798
799
800
801
802
        // Verify it's removed
        let req = Request::builder()
            .method("GET")
            .uri("/list_workers")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        let workers = body_json["urls"].as_array().unwrap();
        assert!(workers.is_empty());
803

804
        ctx.shutdown().await;
805
806
    }

807
808
809
810
    #[tokio::test]
    async fn test_add_worker_invalid_url() {
        let ctx = TestContext::new(vec![]).await;
        let app = ctx.create_app().await;
811

812
813
814
815
816
817
        // Invalid URL format
        let req = Request::builder()
            .method("POST")
            .uri("/add_worker?url=not-a-valid-url")
            .body(Body::empty())
            .unwrap();
818

819
820
        let resp = app.clone().oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
821

822
823
824
825
826
827
        // Missing URL parameter
        let req = Request::builder()
            .method("POST")
            .uri("/add_worker")
            .body(Body::empty())
            .unwrap();
828

829
830
        let resp = app.clone().oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
831

832
833
834
835
836
837
        // Empty URL
        let req = Request::builder()
            .method("POST")
            .uri("/add_worker?url=")
            .body(Body::empty())
            .unwrap();
838

839
840
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
841

842
        ctx.shutdown().await;
843
844
    }

845
846
847
848
849
850
851
852
853
854
855
    #[tokio::test]
    async fn test_add_duplicate_worker() {
        // Start a mock worker
        let mut worker = MockWorker::new(MockWorkerConfig {
            port: 18303,
            worker_type: WorkerType::Regular,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 0.0,
        });
        let url = worker.start().await.unwrap();
856

857
858
        let ctx = TestContext::new(vec![]).await;
        let app = ctx.create_app().await;
859

860
861
862
        // Add worker first time
        let req = Request::builder()
            .method("POST")
863
            .uri(format!("/add_worker?url={}", url))
864
865
866
867
            .body(Body::empty())
            .unwrap();
        let resp = app.clone().oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
868

869
        tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
870

871
872
873
        // Try to add same worker again
        let req = Request::builder()
            .method("POST")
874
            .uri(format!("/add_worker?url={}", url))
875
876
877
878
879
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        // Should return error for duplicate
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
880

881
882
883
        worker.stop().await;
        ctx.shutdown().await;
    }
884

885
886
887
888
889
890
891
892
893
894
895
    #[tokio::test]
    async fn test_add_unhealthy_worker() {
        // Start unhealthy worker
        let mut worker = MockWorker::new(MockWorkerConfig {
            port: 18304,
            worker_type: WorkerType::Regular,
            health_status: HealthStatus::Unhealthy,
            response_delay_ms: 0,
            fail_rate: 0.0,
        });
        let url = worker.start().await.unwrap();
896

897
898
        let ctx = TestContext::new(vec![]).await;
        let app = ctx.create_app().await;
899

900
901
902
        // Try to add unhealthy worker
        let req = Request::builder()
            .method("POST")
903
            .uri(format!("/add_worker?url={}", url))
904
905
906
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
907

908
909
910
911
912
        // Router should reject unhealthy workers
        assert!(
            resp.status() == StatusCode::BAD_REQUEST
                || resp.status() == StatusCode::SERVICE_UNAVAILABLE
        );
913

914
915
        worker.stop().await;
        ctx.shutdown().await;
916
    }
917
918
919
920
921
}

#[cfg(test)]
mod router_policy_tests {
    use super::*;
922

923
924
925
926
927
    #[tokio::test]
    async fn test_random_policy() {
        let ctx = TestContext::new(vec![
            MockWorkerConfig {
                port: 18801,
928
929
930
931
                worker_type: WorkerType::Regular,
                health_status: HealthStatus::Healthy,
                response_delay_ms: 0,
                fail_rate: 0.0,
932
933
934
            },
            MockWorkerConfig {
                port: 18802,
935
                worker_type: WorkerType::Regular,
936
                health_status: HealthStatus::Healthy,
937
938
                response_delay_ms: 0,
                fail_rate: 0.0,
939
940
941
942
943
944
945
946
947
948
949
            },
        ])
        .await;

        // Send multiple requests and verify they succeed
        let app = ctx.create_app().await;

        for i in 0..10 {
            let payload = json!({
                "text": format!("Request {}", i),
                "stream": false
950
951
            });

952
953
954
955
956
957
            let req = Request::builder()
                .method("POST")
                .uri("/generate")
                .header(CONTENT_TYPE, "application/json")
                .body(Body::from(serde_json::to_string(&payload).unwrap()))
                .unwrap();
958

959
960
961
            let resp = app.clone().oneshot(req).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK);
        }
962

963
964
        ctx.shutdown().await;
    }
965

966
967
968
969
970
971
972
973
974
975
976
977
978
979
    #[tokio::test]
    async fn test_worker_selection() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18203,
            worker_type: WorkerType::Regular,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 0.0,
        }])
        .await;

        let _payload = json!({
            "text": "Test selection",
            "stream": false
980
        });
981
982
983
984
985
986
987

        // Check that router has the worker
        let worker_urls = ctx.router.get_worker_urls();
        assert_eq!(worker_urls.len(), 1);
        assert!(worker_urls[0].contains("18203"));

        ctx.shutdown().await;
988
989
990
991
992
993
994
    }
}

#[cfg(test)]
mod error_tests {
    use super::*;

995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
    #[tokio::test]
    async fn test_404_not_found() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18401,
            worker_type: WorkerType::Regular,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 0.0,
        }])
        .await;

        let app = ctx.create_app().await;

        // Test unknown endpoint
        let req = Request::builder()
            .method("GET")
            .uri("/unknown_endpoint")
            .body(Body::empty())
            .unwrap();
1014

1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
        let resp = app.clone().oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::NOT_FOUND);

        // Test POST to unknown endpoint
        let req = Request::builder()
            .method("POST")
            .uri("/api/v2/generate")
            .header(CONTENT_TYPE, "application/json")
            .body(Body::from(
                serde_json::to_string(&json!({"text": "test"})).unwrap(),
            ))
            .unwrap();
1027

1028
1029
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::NOT_FOUND);
1030

1031
1032
        ctx.shutdown().await;
    }
1033

1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
    #[tokio::test]
    async fn test_method_not_allowed() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18402,
            worker_type: WorkerType::Regular,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 0.0,
        }])
        .await;

        let app = ctx.create_app().await;

        // GET request to POST-only endpoint
        let req = Request::builder()
            .method("GET")
            .uri("/generate")
            .body(Body::empty())
            .unwrap();
1053

1054
1055
1056
        let resp = app.clone().oneshot(req).await.unwrap();
        // Note: Axum returns 405 for wrong methods on matched routes
        assert_eq!(resp.status(), StatusCode::METHOD_NOT_ALLOWED);
1057

1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
        // POST request to GET-only endpoint
        let req = Request::builder()
            .method("POST")
            .uri("/health")
            .header(CONTENT_TYPE, "application/json")
            .body(Body::from("{}"))
            .unwrap();

        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::METHOD_NOT_ALLOWED);

        ctx.shutdown().await;
1070
1071
    }

1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
    #[tokio::test]
    async fn test_payload_too_large() {
        // Create context with small payload limit
        let config = RouterConfig {
            mode: RoutingMode::Regular {
                worker_urls: vec![],
            },
            policy: PolicyConfig::Random,
            host: "127.0.0.1".to_string(),
            port: 3010,
            max_payload_size: 1024, // 1KB limit
            request_timeout_secs: 600,
            worker_startup_timeout_secs: 1,
            worker_startup_check_interval_secs: 1,
            dp_aware: false,
            api_key: None,
            discovery: None,
            metrics: None,
            log_dir: None,
            log_level: None,
            request_id_headers: None,
            max_concurrent_requests: 64,
1094
1095
1096
            queue_size: 0,
            queue_timeout_secs: 60,
            rate_limit_tokens_per_second: None,
1097
            cors_allowed_origins: vec![],
1098
            retry: RetryConfig::default(),
1099
            circuit_breaker: CircuitBreakerConfig::default(),
1100
1101
            disable_retries: false,
            disable_circuit_breaker: false,
1102
            health_check: sglang_router_rs::config::HealthCheckConfig::default(),
1103
            enable_igw: false,
1104
1105
1106
1107
1108
1109
        };

        let ctx = TestContext::new_with_config(
            config,
            vec![MockWorkerConfig {
                port: 18403,
1110
1111
1112
1113
                worker_type: WorkerType::Regular,
                health_status: HealthStatus::Healthy,
                response_delay_ms: 0,
                fail_rate: 0.0,
1114
1115
1116
            }],
        )
        .await;
1117

1118
1119
1120
        // Note: The server would have payload size middleware configured
        // but we cannot test it directly through the test app
        // This test is kept for documentation purposes
1121

1122
1123
        ctx.shutdown().await;
    }
1124

1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
    #[tokio::test]
    async fn test_invalid_json_payload() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18404,
            worker_type: WorkerType::Regular,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 0.0,
        }])
        .await;

        let app = ctx.create_app().await;

        // Send invalid JSON
        let req = Request::builder()
            .method("POST")
            .uri("/generate")
            .header(CONTENT_TYPE, "application/json")
            .body(Body::from("{invalid json}"))
            .unwrap();
1145

1146
1147
        let resp = app.clone().oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
1148

1149
1150
1151
1152
1153
1154
1155
        // Send empty body
        let req = Request::builder()
            .method("POST")
            .uri("/generate")
            .header(CONTENT_TYPE, "application/json")
            .body(Body::empty())
            .unwrap();
1156

1157
1158
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
1159

1160
1161
        ctx.shutdown().await;
    }
1162

1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
    #[tokio::test]
    async fn test_missing_required_fields() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18405,
            worker_type: WorkerType::Regular,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 0.0,
        }])
        .await;

        let app = ctx.create_app().await;

        // Missing messages in chat completion
        let payload = json!({
            "model": "test-model"
            // missing "messages"
1180
        });
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193

        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();
        // Axum validates JSON schema - returns 422 for validation errors
        assert_eq!(resp.status(), StatusCode::UNPROCESSABLE_ENTITY);

        ctx.shutdown().await;
1194
1195
    }

1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
    #[tokio::test]
    async fn test_invalid_model() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18406,
            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": "invalid-model-name-that-does-not-exist",
            "messages": [{"role": "user", "content": "Hello"}],
            "stream": false
        });
1214

1215
1216
1217
1218
1219
1220
        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();
1221

1222
1223
1224
        let resp = app.oneshot(req).await.unwrap();
        // Mock worker accepts any model, but real implementation might return 400
        assert!(resp.status().is_success() || resp.status() == StatusCode::BAD_REQUEST);
1225

1226
1227
1228
        ctx.shutdown().await;
    }
}
1229

1230
1231
1232
#[cfg(test)]
mod cache_tests {
    use super::*;
1233

1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
    #[tokio::test]
    async fn test_flush_cache() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18501,
            worker_type: WorkerType::Regular,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 0.0,
        }])
        .await;

        let app = ctx.create_app().await;

        let req = Request::builder()
            .method("POST")
            .uri("/flush_cache")
            .body(Body::empty())
            .unwrap();
1252

1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);

        // The response might be empty or contain a message
        let body_bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        if !body_bytes.is_empty() {
            if let Ok(body) = serde_json::from_slice::<serde_json::Value>(&body_bytes) {
                // Check that we got a successful response with expected fields
                assert!(body.is_object());
                assert!(body.get("message").is_some() || body.get("status").is_some());
            }
        }

        ctx.shutdown().await;
1269
1270
    }

1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
    #[tokio::test]
    async fn test_get_loads() {
        let ctx = TestContext::new(vec![
            MockWorkerConfig {
                port: 18502,
                worker_type: WorkerType::Regular,
                health_status: HealthStatus::Healthy,
                response_delay_ms: 0,
                fail_rate: 0.0,
            },
            MockWorkerConfig {
                port: 18503,
1283
1284
1285
1286
                worker_type: WorkerType::Regular,
                health_status: HealthStatus::Healthy,
                response_delay_ms: 0,
                fail_rate: 0.0,
1287
1288
1289
            },
        ])
        .await;
1290

1291
        let app = ctx.create_app().await;
1292

1293
1294
1295
1296
1297
        let req = Request::builder()
            .method("GET")
            .uri("/get_loads")
            .body(Body::empty())
            .unwrap();
1298

1299
1300
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
1301

1302
1303
1304
1305
        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
1306

1307
1308
1309
1310
        // Verify the response contains load information
        assert!(body_json.is_object());
        // The exact structure depends on the implementation
        // but should contain worker load information
1311

1312
1313
        ctx.shutdown().await;
    }
1314

1315
1316
1317
    #[tokio::test]
    async fn test_flush_cache_no_workers() {
        let ctx = TestContext::new(vec![]).await;
1318

1319
        let app = ctx.create_app().await;
1320

1321
1322
1323
1324
1325
        let req = Request::builder()
            .method("POST")
            .uri("/flush_cache")
            .body(Body::empty())
            .unwrap();
1326

1327
1328
1329
1330
1331
        let resp = app.oneshot(req).await.unwrap();
        // Should either succeed (no-op) or return service unavailable
        assert!(
            resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE
        );
1332

1333
        ctx.shutdown().await;
1334
1335
1336
1337
    }
}

#[cfg(test)]
1338
mod load_balancing_tests {
1339
1340
    use super::*;

1341
1342
1343
1344
1345
1346
    #[tokio::test]
    async fn test_request_distribution() {
        // Create multiple workers
        let ctx = TestContext::new(vec![
            MockWorkerConfig {
                port: 18601,
1347
1348
1349
1350
                worker_type: WorkerType::Regular,
                health_status: HealthStatus::Healthy,
                response_delay_ms: 0,
                fail_rate: 0.0,
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
            },
            MockWorkerConfig {
                port: 18602,
                worker_type: WorkerType::Regular,
                health_status: HealthStatus::Healthy,
                response_delay_ms: 0,
                fail_rate: 0.0,
            },
        ])
        .await;
1361

1362
        let app = ctx.create_app().await;
1363

1364
1365
1366
1367
1368
1369
1370
        // Send multiple requests and track distribution
        let mut request_count = 0;
        for i in 0..10 {
            let payload = json!({
                "text": format!("Request {}", i),
                "stream": false
            });
1371

1372
1373
1374
1375
1376
1377
            let req = Request::builder()
                .method("POST")
                .uri("/generate")
                .header(CONTENT_TYPE, "application/json")
                .body(Body::from(serde_json::to_string(&payload).unwrap()))
                .unwrap();
1378

1379
1380
1381
            let resp = app.clone().oneshot(req).await.unwrap();
            if resp.status() == StatusCode::OK {
                request_count += 1;
1382
            }
1383
        }
1384

1385
1386
        // With random policy, all requests should succeed
        assert_eq!(request_count, 10);
1387

1388
1389
1390
        ctx.shutdown().await;
    }
}
1391

1392
1393
1394
#[cfg(test)]
mod pd_mode_tests {
    use super::*;
1395

1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
    #[tokio::test]
    async fn test_pd_mode_routing() {
        // Create PD mode configuration with prefill and decode workers
        let mut prefill_worker = MockWorker::new(MockWorkerConfig {
            port: 18701,
            worker_type: WorkerType::Prefill,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 0.0,
        });
1406

1407
1408
1409
1410
1411
1412
        let mut decode_worker = MockWorker::new(MockWorkerConfig {
            port: 18702,
            worker_type: WorkerType::Decode,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 0.0,
1413
1414
        });

1415
1416
        let prefill_url = prefill_worker.start().await.unwrap();
        let decode_url = decode_worker.start().await.unwrap();
1417

1418
1419
1420
1421
1422
        tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;

        // Extract port from prefill URL
        let prefill_port = prefill_url
            .split(':')
1423
            .next_back()
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
            .and_then(|p| p.trim_end_matches('/').parse::<u16>().ok())
            .unwrap_or(9000);

        let config = RouterConfig {
            mode: RoutingMode::PrefillDecode {
                prefill_urls: vec![(prefill_url, Some(prefill_port))],
                decode_urls: vec![decode_url],
                prefill_policy: None,
                decode_policy: None,
            },
            policy: PolicyConfig::Random,
            host: "127.0.0.1".to_string(),
            port: 3011,
            max_payload_size: 256 * 1024 * 1024,
            request_timeout_secs: 600,
            worker_startup_timeout_secs: 1,
            worker_startup_check_interval_secs: 1,
            discovery: None,
            metrics: None,
            log_dir: None,
            dp_aware: false,
            api_key: None,
            log_level: None,
            request_id_headers: None,
            max_concurrent_requests: 64,
1449
1450
1451
            queue_size: 0,
            queue_timeout_secs: 60,
            rate_limit_tokens_per_second: None,
1452
            cors_allowed_origins: vec![],
1453
            retry: RetryConfig::default(),
1454
            circuit_breaker: CircuitBreakerConfig::default(),
1455
1456
            disable_retries: false,
            disable_circuit_breaker: false,
1457
            health_check: sglang_router_rs::config::HealthCheckConfig::default(),
1458
            enable_igw: false,
1459
1460
        };

1461
1462
1463
        // Create app context
        let app_context = common::create_test_context(config);

1464
        // Create router - this might fail due to health check issues
1465
        let router_result = RouterFactory::create_router(&app_context).await;
1466
1467
1468
1469
1470
1471
1472

        // Clean up workers
        prefill_worker.stop().await;
        decode_worker.stop().await;

        // For now, just verify the configuration was attempted
        assert!(router_result.is_err() || router_result.is_ok());
1473
1474
1475
1476
    }
}

#[cfg(test)]
1477
mod request_id_tests {
1478
1479
    use super::*;

1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
    #[tokio::test]
    async fn test_request_id_generation() {
        let ctx = TestContext::new(vec![MockWorkerConfig {
            port: 18901,
            worker_type: WorkerType::Regular,
            health_status: HealthStatus::Healthy,
            response_delay_ms: 0,
            fail_rate: 0.0,
        }])
        .await;

        let app = ctx.create_app().await;

        // Test 1: Request without any request ID header should generate one
        let payload = json!({
            "text": "Test request",
            "stream": false
        });
1498

1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
        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();
        assert_eq!(resp.status(), StatusCode::OK);

        // Check that response has x-request-id header
        let request_id = resp.headers().get("x-request-id");
        assert!(
            request_id.is_some(),
            "Response should have x-request-id header"
        );

        let id_value = request_id.unwrap().to_str().unwrap();
        assert!(
            id_value.starts_with("gnt-"),
            "Generate endpoint should have gnt- prefix"
        );
        assert!(
            id_value.len() > 4,
            "Request ID should have content after prefix"
        );

        // Test 2: Request with custom x-request-id should preserve it
        let custom_id = "custom-request-id-123";
        let req = Request::builder()
            .method("POST")
            .uri("/generate")
            .header(CONTENT_TYPE, "application/json")
            .header("x-request-id", custom_id)
            .body(Body::from(serde_json::to_string(&payload).unwrap()))
            .unwrap();

        let resp = app.clone().oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
1538

1539
1540
1541
1542
1543
1544
1545
1546
        let response_id = resp.headers().get("x-request-id");
        assert!(response_id.is_some());
        assert_eq!(response_id.unwrap(), custom_id);

        // Test 3: Different endpoints should have different prefixes
        let chat_payload = json!({
            "messages": [{"role": "user", "content": "Hello"}],
            "model": "test-model"
1547
        });
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583

        let req = Request::builder()
            .method("POST")
            .uri("/v1/chat/completions")
            .header(CONTENT_TYPE, "application/json")
            .body(Body::from(serde_json::to_string(&chat_payload).unwrap()))
            .unwrap();

        let resp = app.clone().oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);

        let request_id = resp.headers().get("x-request-id");
        assert!(request_id.is_some());
        assert!(request_id
            .unwrap()
            .to_str()
            .unwrap()
            .starts_with("chatcmpl-"));

        // Test 4: Alternative request ID headers should be recognized
        let req = Request::builder()
            .method("POST")
            .uri("/generate")
            .header(CONTENT_TYPE, "application/json")
            .header("x-correlation-id", "correlation-123")
            .body(Body::from(serde_json::to_string(&payload).unwrap()))
            .unwrap();

        let resp = app.clone().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(), "correlation-123");

        ctx.shutdown().await;
1584
1585
    }

1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
    #[tokio::test]
    async fn test_request_id_with_custom_headers() {
        // Create config with custom request ID headers
        let config = RouterConfig {
            mode: RoutingMode::Regular {
                worker_urls: vec![],
            },
            policy: PolicyConfig::Random,
            host: "127.0.0.1".to_string(),
            port: 3002,
            max_payload_size: 256 * 1024 * 1024,
            request_timeout_secs: 600,
            worker_startup_timeout_secs: 1,
            worker_startup_check_interval_secs: 1,
            discovery: None,
            metrics: None,
            dp_aware: false,
            api_key: None,
            log_dir: None,
            log_level: None,
            request_id_headers: Some(vec!["custom-id".to_string(), "trace-id".to_string()]),
            max_concurrent_requests: 64,
1608
1609
1610
            queue_size: 0,
            queue_timeout_secs: 60,
            rate_limit_tokens_per_second: None,
1611
            cors_allowed_origins: vec![],
1612
            retry: RetryConfig::default(),
1613
            circuit_breaker: CircuitBreakerConfig::default(),
1614
1615
            disable_retries: false,
            disable_circuit_breaker: false,
1616
            health_check: sglang_router_rs::config::HealthCheckConfig::default(),
1617
            enable_igw: false,
1618
        };
1619

1620
1621
1622
1623
1624
        let ctx = TestContext::new_with_config(
            config,
            vec![MockWorkerConfig {
                port: 18902,
                worker_type: WorkerType::Regular,
1625
1626
1627
                health_status: HealthStatus::Healthy,
                response_delay_ms: 0,
                fail_rate: 0.0,
1628
1629
1630
            }],
        )
        .await;
1631

1632
        let app = ctx.create_app().await;
1633

1634
1635
1636
1637
        let payload = json!({
            "text": "Test request",
            "stream": false
        });
1638

1639
1640
1641
1642
1643
1644
1645
1646
        // Test custom header is recognized
        let req = Request::builder()
            .method("POST")
            .uri("/generate")
            .header(CONTENT_TYPE, "application/json")
            .header("custom-id", "my-custom-id")
            .body(Body::from(serde_json::to_string(&payload).unwrap()))
            .unwrap();
1647

1648
1649
1650
1651
1652
1653
1654
1655
        let resp = app.clone().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(), "my-custom-id");

        ctx.shutdown().await;
1656
1657
    }
}