fix precommit unit test failures

Signed-off-by: Brandon Pelfrey <bpelfrey@nvidia.com>

Author: brandonpelfrey

tests/unit/common/config/test_user_config.py

@@ -39,6 +39,9 @@ def make_endpoint(
     **kwargs,
 ) -> EndpointConfig:
     """Create an EndpointConfig with sensible defaults for testing."""
+    # Convert url= to urls= for backward compatibility
+    if "url" in kwargs:
+        kwargs["urls"] = [kwargs.pop("url")]
     return EndpointConfig(
         model_names=model_names or ["test-model"],
         type=endpoint_type,
@@ -97,7 +100,7 @@ def test_to_json_string(self):
                 type=EndpointType.CHAT,
                 custom_endpoint="custom_endpoint",
                 streaming=True,
-                url="http://custom-url",
+                urls=["http://custom-url"],
                 api_key="test_api_key",
                 timeout_seconds=10,
             ),

tests/unit/server_metrics/test_server_metrics_manager.py

@@ -21,7 +21,7 @@ def user_config_with_endpoint() -> UserConfig:
         endpoint=EndpointConfig(
             model_names=["test-model"],
             type=EndpointType.CHAT,
-            url="http://localhost:8000/v1/chat",
+            urls=["http://localhost:8000/v1/chat"],
         ),
     )
 
@@ -33,7 +33,7 @@ def user_config_with_server_metrics_urls() -> UserConfig:
         endpoint=EndpointConfig(
             model_names=["test-model"],
             type=EndpointType.CHAT,
-            url="http://localhost:8000/v1/chat",
+            urls=["http://localhost:8000/v1/chat"],
         ),
         server_metrics=[
             "http://custom-endpoint:9400/metrics",
@@ -423,7 +423,7 @@ async def test_configure_when_server_metrics_disabled(
             endpoint=EndpointConfig(
                 model_names=["test-model"],
                 type=EndpointType.CHAT,
-                url="http://localhost:8000/v1/chat",
+                urls=["http://localhost:8000/v1/chat"],
             ),
             no_server_metrics=True,  # Disable server metrics
         )

tests/unit/timing/test_timing_manager.py

@@ -24,8 +24,11 @@
 
 @pytest.fixture
 def user_config() -> UserConfig:
+    mock_endpoint = MagicMock()
+    mock_endpoint.urls = ["http://localhost:8000"]
+    mock_endpoint.url_selection_strategy = "round_robin"
     return UserConfig.model_construct(
-        endpoint=MagicMock(), _timing_mode=TimingMode.REQUEST_RATE
+        endpoint=mock_endpoint, _timing_mode=TimingMode.REQUEST_RATE
     )
 
 
@@ -62,6 +65,14 @@ def mock_metadata() -> DatasetMetadata:
     )
 
 
+def _create_mock_endpoint() -> MagicMock:
+    """Create a mock endpoint with required URL attributes."""
+    mock_endpoint = MagicMock()
+    mock_endpoint.urls = ["http://localhost:8000"]
+    mock_endpoint.url_selection_strategy = "round_robin"
+    return mock_endpoint
+
+
 class TestTimingManagerDatasetConfiguration:
     @pytest.mark.parametrize(
         "timing_mode", [TimingMode.FIXED_SCHEDULE, TimingMode.REQUEST_RATE]
@@ -70,7 +81,9 @@ class TestTimingManagerDatasetConfiguration:
     async def test_profile_configure_waits_for_dataset_notification(
         self, create_manager, mock_metadata, timing_mode
     ) -> None:
-        cfg = UserConfig.model_construct(endpoint=MagicMock(), _timing_mode=timing_mode)
+        cfg = UserConfig.model_construct(
+            endpoint=_create_mock_endpoint(), _timing_mode=timing_mode
+        )
         mgr = create_manager(cfg)
         mock_engine = MagicMock()
         mock_engine.initialize = lambda *a, **kw: asyncio.sleep(0)
@@ -105,7 +118,7 @@ async def test_profile_configure_waits_for_dataset_notification(
     @pytest.mark.asyncio
     async def test_dataset_configuration_timeout(self, create_manager) -> None:
         cfg = UserConfig.model_construct(
-            endpoint=MagicMock(), _timing_mode=TimingMode.FIXED_SCHEDULE
+            endpoint=_create_mock_endpoint(), _timing_mode=TimingMode.FIXED_SCHEDULE
         )
         mgr = create_manager(cfg)
         with (
@@ -123,7 +136,7 @@ async def test_dataset_notification_before_configure(
         self, create_manager, mock_metadata
     ) -> None:
         cfg = UserConfig.model_construct(
-            endpoint=MagicMock(), _timing_mode=TimingMode.FIXED_SCHEDULE
+            endpoint=_create_mock_endpoint(), _timing_mode=TimingMode.FIXED_SCHEDULE
         )
         mgr = create_manager(cfg)
         await mgr._on_dataset_configured_notification(

tests/unit/transports/conftest.py

@@ -54,7 +54,7 @@ def create_model_endpoint_info(
         ),
         endpoint=EndpointInfo(
             type=EndpointType.CHAT,
-            base_url=base_url,
+            base_urls=[base_url],
             custom_endpoint=custom_endpoint,
             streaming=streaming,
             api_key=api_key,
@@ -114,7 +114,7 @@ def user_config() -> UserConfig:
     return UserConfig(
         endpoint=EndpointConfig(
             type=EndpointType.CHAT,
-            url="http://localhost:8000",
+            urls=["http://localhost:8000"],
             timeout_seconds=600,
             model_names=["gpt-4"],
             api_key="test-api-key",

tests/unit/transports/test_aiohttp_transport.py

@@ -153,15 +153,14 @@ def test_get_transport_headers(self, transport, streaming, expected_accept):
         ],
         ids=["http-prefix", "no-scheme", "https-prefix"],
     )
-    def test_get_url(
-        self, model_endpoint_non_streaming, base_url, custom_endpoint, expected_url
-    ):
+    def test_get_url(self, base_url, custom_endpoint, expected_url):
         """Test get_url with various base URLs and endpoints."""
-        model_endpoint_non_streaming.endpoint.base_url = base_url
-        model_endpoint_non_streaming.endpoint.custom_endpoint = custom_endpoint
+        model_endpoint = create_model_endpoint_info(
+            base_url=base_url, custom_endpoint=custom_endpoint
+        )
 
-        transport = AioHttpTransport(model_endpoint=model_endpoint_non_streaming)
-        request_info = create_request_info(model_endpoint_non_streaming)
+        transport = AioHttpTransport(model_endpoint=model_endpoint)
+        request_info = create_request_info(model_endpoint)
         url = transport.get_url(request_info)
         assert url == expected_url
 
@@ -458,7 +457,7 @@ async def test_full_request_flow_non_streaming(self):
             "headers": transport.aiohttp_client.post_request.call_args[0][2],
         }
 
-        assert "https://api.example.com/v1/chat/completions" in args["url"]
+        assert args["url"].startswith("https://api.example.com/v1/chat/completions")
         assert "api-version=2024-10-01" in args["url"]
         assert "Hello" in args["json_str"]
         assert args["headers"]["Authorization"] == "Bearer test-key"

tests/unit/transports/test_base_transport.py

@@ -24,6 +24,13 @@
 AIPERF_USER_AGENT = f"aiperf/{importlib_metadata.version('aiperf')}"
 
 
+def _ensure_scheme(url: str) -> str:
+    """Ensure URL has a scheme prefix."""
+    if not url.startswith(("http://", "https://")):
+        return f"http://{url}"
+    return url
+
+
 class FakeTransport(BaseTransport):
     """Concrete implementation of BaseTransport for testing."""
 
@@ -36,7 +43,9 @@ def metadata(cls) -> TransportMetadata:
 
     def get_url(self, request_info: RequestInfo) -> str:
         endpoint_info = request_info.model_endpoint.endpoint
-        base_url = endpoint_info.base_url or ""
+        base_url = (
+            _ensure_scheme(endpoint_info.base_url) if endpoint_info.base_url else ""
+        )
         if endpoint_info.custom_endpoint:
             return f"{base_url}{endpoint_info.custom_endpoint}"
         return base_url
@@ -60,7 +69,7 @@ def model_endpoint(self):
             ),
             endpoint=EndpointInfo(
                 type=EndpointType.CHAT,
-                base_url="http://localhost:8000",
+                base_urls=["http://localhost:8000"],
                 custom_endpoint="/v1/chat/completions",
             ),
         )
@@ -200,11 +209,20 @@ def test_build_url_with_endpoint_params(self, transport, request_info):
         assert "timeout=30" in url
         assert url.startswith("http://localhost:8000/v1/chat/completions?")
 
-    def test_build_url_preserves_existing_params(self, transport, model_endpoint):
+    def test_build_url_preserves_existing_params(self, transport):
         """Test that existing URL params are preserved."""
-        model_endpoint.endpoint.base_url = (
-            "http://localhost:8000/v1/chat/completions?existing=param"
+        model_endpoint = ModelEndpointInfo(
+            models=ModelListInfo(
+                models=[ModelInfo(name="test-model")],
+                model_selection_strategy=ModelSelectionStrategy.ROUND_ROBIN,
+            ),
+            endpoint=EndpointInfo(
+                type=EndpointType.CHAT,
+                base_urls=["http://localhost:8000/v1/chat/completions?existing=param"],
+                custom_endpoint=None,
+            ),
         )
+        transport = FakeTransport(model_endpoint=model_endpoint)
         request_info = RequestInfo(
             model_endpoint=model_endpoint,
             turns=[],
@@ -222,13 +240,20 @@ def test_build_url_preserves_existing_params(self, transport, model_endpoint):
         assert "existing=param" in url
         assert "new=value" in url
 
-    def test_build_url_endpoint_params_override_existing(
-        self, transport, model_endpoint
-    ):
+    def test_build_url_endpoint_params_override_existing(self, transport):
         """Test that endpoint params override existing URL params."""
-        model_endpoint.endpoint.base_url = (
-            "http://localhost:8000/v1/chat/completions?key=original"
+        model_endpoint = ModelEndpointInfo(
+            models=ModelListInfo(
+                models=[ModelInfo(name="test-model")],
+                model_selection_strategy=ModelSelectionStrategy.ROUND_ROBIN,
+            ),
+            endpoint=EndpointInfo(
+                type=EndpointType.CHAT,
+                base_urls=["http://localhost:8000/v1/chat/completions?key=original"],
+                custom_endpoint=None,
+            ),
         )
+        transport = FakeTransport(model_endpoint=model_endpoint)
         request_info = RequestInfo(
             model_endpoint=model_endpoint,
             turns=[],
@@ -267,9 +292,20 @@ def test_build_url_no_params_preserves_clean_url(self, transport, request_info):
         assert "?" not in url
         assert url == "http://localhost:8000/v1/chat/completions"
 
-    def test_build_url_complex_query_string(self, transport, model_endpoint):
+    def test_build_url_complex_query_string(self, transport):
         """Test complex query string handling."""
-        model_endpoint.endpoint.base_url = "http://localhost:8000/api?a=1&b=2&c=3"
+        model_endpoint = ModelEndpointInfo(
+            models=ModelListInfo(
+                models=[ModelInfo(name="test-model")],
+                model_selection_strategy=ModelSelectionStrategy.ROUND_ROBIN,
+            ),
+            endpoint=EndpointInfo(
+                type=EndpointType.CHAT,
+                base_urls=["http://localhost:8000/api?a=1&b=2&c=3"],
+                custom_endpoint=None,
+            ),
+        )
+        transport = FakeTransport(model_endpoint=model_endpoint)
         request_info = RequestInfo(
             model_endpoint=model_endpoint,
             turns=[],