test: test_busy_threshold_endpoint on both NATS and TCP making sure info for rejection is working

Signed-off-by: Jacky <[email protected]>

Author: kthui

tests/router/common.py

@@ -4,6 +4,7 @@
 import asyncio
 import json
 import logging
+import os
 import random
 import string
 import time
@@ -38,6 +39,7 @@ def __init__(
         store_backend: str = "etcd",
         enforce_disagg: bool = False,
         busy_threshold: float | None = None,
+        request_plane: str = "nats",
     ):
         command = [
             "python3",
@@ -61,8 +63,12 @@ def __init__(
         if busy_threshold is not None:
             command.extend(["--busy-threshold", str(busy_threshold)])
 
+        env = os.environ.copy()
+        env["DYN_REQUEST_PLANE"] = request_plane
+
         super().__init__(
             command=command,
+            env=env,
             timeout=60,
             display_output=True,
             health_check_ports=[frontend_port],
@@ -1895,6 +1901,7 @@ def _test_busy_threshold_endpoint(
     frontend_port: int,
     test_payload: dict,
     store_backend: str = "etcd",
+    request_plane: str = "nats",
 ):
     """Test that the /busy_threshold endpoint can be hit and responds correctly.
 
@@ -1912,6 +1919,7 @@ def _test_busy_threshold_endpoint(
         frontend_port: Port for the frontend HTTP server
         test_payload: Base test payload (used to extract model name)
         store_backend: Storage backend to use ("etcd" or "file"). Defaults to "etcd".
+        request_plane: Request plane to use ("nats" or "tcp"). Defaults to "nats".
 
     Raises:
         AssertionError: If endpoint responses are incorrect
@@ -1929,6 +1937,7 @@ def _test_busy_threshold_endpoint(
             engine_workers.namespace,
             store_backend,
             busy_threshold=initial_threshold,
+            request_plane=request_plane,
         )
         kv_router.__enter__()
 

tests/router/test_router_e2e_with_mockers.py

@@ -175,6 +175,7 @@ def __init__(
         mocker_args: Optional[Dict[str, Any]] = None,
         num_mockers: int = 1,
         store_backend: str = "etcd",
+        request_plane: str = "nats",
     ):
         namespace_suffix = generate_random_suffix()
         self.namespace = f"test-namespace-{namespace_suffix}"
@@ -191,8 +192,12 @@ def __init__(
             mocker_args=mocker_args,
         )
 
+        env = os.environ.copy()
+        env["DYN_REQUEST_PLANE"] = request_plane
+
         self._process = ManagedProcess(
             command=command,
+            env=env,
             timeout=60,
             display_output=True,
             health_check_ports=[],
@@ -649,8 +654,9 @@ def test_router_disagg_decisions(
 
 
 @pytest.mark.parallel
+@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
 def test_busy_threshold_endpoint(
-    request, runtime_services_session, predownload_tokenizers
+    request, runtime_services_session, predownload_tokenizers, request_plane
 ):
     """Test that the /busy_threshold endpoint can be hit and responds correctly.
 
@@ -661,14 +667,19 @@ def test_busy_threshold_endpoint(
 
     For now, this test only verifies the endpoint is accessible and returns valid responses.
     """
-    logger.info("Starting busy_threshold endpoint test")
+    logger.info(
+        f"Starting busy_threshold endpoint test with request_plane={request_plane}"
+    )
 
     mocker_args = {"speedup_ratio": SPEEDUP_RATIO, "block_size": BLOCK_SIZE}
 
     try:
         logger.info(f"Starting {NUM_MOCKERS} mocker instances")
         mockers = MockerProcess(
-            request, mocker_args=mocker_args, num_mockers=NUM_MOCKERS
+            request,
+            mocker_args=mocker_args,
+            num_mockers=NUM_MOCKERS,
+            request_plane=request_plane,
         )
         logger.info(f"All mockers using endpoint: {mockers.endpoint}")
         mockers.__enter__()
@@ -681,6 +692,7 @@ def test_busy_threshold_endpoint(
             request=request,
             frontend_port=frontend_port,
             test_payload=TEST_PAYLOAD,
+            request_plane=request_plane,
         )
 
     finally: