From 4a10e37ba7083e34660fd50267a0e32ebba77a80 Mon Sep 17 00:00:00 2001
From: Xinyue Zhang <xinyue.zhang@oracle.com>
Date: Sat, 15 Nov 2025 02:00:25 -0800
Subject: [PATCH] [router] Fix flaky router e2e tests (#13306)

---
 .../integration_mock/test_circuit_breaker.py  | 33 ++++++++++++-------
 .../py_test/integration_mock/test_retries.py  | 32 +++++++++++-------
 2 files changed, 43 insertions(+), 22 deletions(-)

diff --git a/sgl-router/py_test/integration_mock/test_circuit_breaker.py b/sgl-router/py_test/integration_mock/test_circuit_breaker.py
index b6d05d577..01e1213c7 100644
--- a/sgl-router/py_test/integration_mock/test_circuit_breaker.py
+++ b/sgl-router/py_test/integration_mock/test_circuit_breaker.py
@@ -112,17 +112,28 @@ def test_circuit_breaker_disable_flag(router_manager, mock_workers):
             "disable_retries": True,
         },
     )
-    r = requests.post(
-        f"{rh.url}/v1/completions",
-        json={
-            "model": "test-model",
-            "prompt": "x",
-            "max_tokens": 1,
-            "stream": False,
-        },
-        timeout=3,
-    )
-    assert r.status_code == 500
+
+    saw_500 = False
+    for _ in range(8):
+        r = requests.post(
+            f"{rh.url}/v1/completions",
+            json={
+                "model": "test-model",
+                "prompt": "x",
+                "max_tokens": 1,
+                "stream": False,
+            },
+            timeout=3,
+        )
+        if r.status_code == 500:
+            # Worker starts, continue to check
+            saw_500 = True
+            break
+        assert (
+            r.status_code == 503
+        ), "Should only see 503 when waiting for worker to start"
+
+    assert saw_500
 
 
 @pytest.mark.integration
diff --git a/sgl-router/py_test/integration_mock/test_retries.py b/sgl-router/py_test/integration_mock/test_retries.py
index c5f61049b..0c88ca7d6 100644
--- a/sgl-router/py_test/integration_mock/test_retries.py
+++ b/sgl-router/py_test/integration_mock/test_retries.py
@@ -47,15 +47,25 @@ def test_disable_retries_surfaces_failure(router_manager, mock_workers):
         },
     )
 
-    r = requests.post(
-        f"{rh.url}/v1/completions",
-        json={
-            "model": "test-model",
-            "prompt": "x",
-            "max_tokens": 1,
-            "stream": False,
-        },
-        timeout=5,
-    )
-    assert r.status_code == 500
+    saw_500 = False
+    for _ in range(8):
+        r = requests.post(
+            f"{rh.url}/v1/completions",
+            json={
+                "model": "test-model",
+                "prompt": "x",
+                "max_tokens": 1,
+                "stream": False,
+            },
+            timeout=5,
+        )
+        if r.status_code == 500:
+            # Worker starts, continue to check
+            saw_500 = True
+            break
+        assert (
+            r.status_code == 503
+        ), "Should only see 503 when waiting for worker to start"
+
+    assert saw_500
     # mock_workers fixture handles cleanup