fix: Swagger UI message print when `disable_openapi_url=False

src/litserve/server.py

@@ -752,6 +752,7 @@ def __init__(
         self.timeout = timeout
         self.litapi_connector.set_request_timeout(timeout)
         self.app = FastAPI(lifespan=self.lifespan, openapi_url="" if disable_openapi_url else "/openapi.json")
+        self._disable_openapi_url = disable_openapi_url
 
         self.app.response_queue_id = None
         self.response_buffer = {}
@@ -1345,7 +1346,8 @@ def run(
             uvicorn_workers = self._start_server(
                 port, num_api_servers, log_level, sockets, api_server_worker_type, **kwargs
             )
-            print(f"Swagger UI is available at http://0.0.0.0:{port}/docs")
+            if not self._disable_openapi_url:
+                print(f"Swagger UI is available at http://0.0.0.0:{port}/docs")
 
             self._start_worker_monitoring(manager, uvicorn_workers)
 

tests/unit/test_lit_server.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import asyncio
+import contextlib
 import json
 import sys
 import time
@@ -338,6 +339,26 @@ def test_server_terminate():
         server._transport.close.assert_called()
 
 
+@pytest.mark.parametrize(("disable_openapi_url", "should_print"), [(False, True), (True, False)])
+@patch("builtins.print")
+@patch("litserve.server.uvicorn")
+def test_disable_openapi_url_print_message(mock_uvicorn, mock_print, mock_manager, disable_openapi_url, should_print):
+    """Test that the Swagger UI message is only printed when disable_openapi_url=False."""
+    server = LitServer(SimpleLitAPI(), disable_openapi_url=disable_openapi_url)
+    server.verify_worker_status = MagicMock()
+
+    with (
+        patch("litserve.server.mp.Manager", return_value=mock_manager),
+        contextlib.suppress(Exception),
+    ):
+        server.run(port=8000)
+
+    if should_print:
+        mock_print.assert_called_with("Swagger UI is available at http://0.0.0.0:8000/docs")
+    else:
+        mock_print.assert_not_called()
+
+
 class IdentityAPI(ls.test_examples.SimpleLitAPI):
     def predict(self, x, context):
         context["input"] = x