| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from argparse import Namespace |
| 5 | from typing import get_args |
| 6 | |
| 7 | import pytest |
| 8 | import regex as re |
| 9 | from fastapi import FastAPI |
| 10 | from starlette.responses import JSONResponse |
| 11 | from starlette.routing import Route |
| 12 | from starlette.testclient import TestClient |
| 13 | |
| 14 | from vllm.entrypoints.launchers.api_server.routers import register_api_routers |
| 15 | from vllm.entrypoints.serve.middleware.authenticate import ( |
| 16 | GUARDED_PREFIX, |
| 17 | AuthenticationMiddleware, |
| 18 | ) |
| 19 | from vllm.tasks import POOLING_TASKS, SupportedTask |
| 20 | |
| 21 | # --------------------------------------------------------------------------- |
| 22 | # Helpers |
| 23 | # --------------------------------------------------------------------------- |
| 24 | |
| 25 | |
| 26 | def get_all_http_routes(app: FastAPI) -> list[tuple[str, list[str]]]: |
| 27 | """Extract all HTTP routes (path, methods) from the FastAPI app.""" |
| 28 | routes = [] |
| 29 | for route in app.routes: |
| 30 | if not isinstance(route, Route): |
| 31 | continue |
| 32 | path = route.path |
| 33 | methods = list(route.methods or {"GET"}) |
| 34 | routes.append((path, methods)) |
| 35 | return routes |
| 36 | |
| 37 | |
| 38 | def generate_test_path(path_template: str) -> str: |
| 39 | """Replace path parameters (e.g. {response_id}) with 'test'.""" |
| 40 | return re.sub(r"\{[^}]+\}", "test", path_template) |
| 41 | |
| 42 | |
| 43 | def _create_app_with_mock_routes(routes: list[tuple[str, list[str]]]) -> FastAPI: |
| 44 | """Create a FastAPI app with AuthenticationMiddleware and mock endpoints.""" |
| 45 | app = FastAPI() |
| 46 | app.add_middleware(AuthenticationMiddleware, tokens=["valid-token"]) |
| 47 | |
| 48 | async def mock_endpoint(): |
| 49 | return JSONResponse({"status": "ok"}) |
| 50 | |
| 51 | for path_template, methods in routes: |
| 52 | allowed_methods = list(set(methods + ["OPTIONS"])) |
| 53 | app.add_api_route( |
| 54 | path_template, |
| 55 | mock_endpoint, |
| 56 | methods=allowed_methods, |
| 57 | include_in_schema=False, |
| 58 | ) |
| 59 | return app |
| 60 | |
| 61 | |
| 62 | class MockModelConfig: |
| 63 | def __init__(self): |
| 64 | self.hf_config = Namespace() |
| 65 | self.hf_config.num_labels = 1 |
| 66 | |
| 67 | def get_pooling_task(self, supported_tasks: tuple["SupportedTask", ...]): |
| 68 | pooling_tasks = [s for s in supported_tasks if s in POOLING_TASKS] |
| 69 | return pooling_tasks[0] if len(pooling_tasks) > 0 else None |
| 70 | |
| 71 | |
| 72 | @pytest.fixture(params=get_args(SupportedTask)) |
| 73 | def task_routes(request, monkeypatch) -> tuple[str, list[tuple[str, list[str]]]]: |
| 74 | """For each supported task, build an app with only that task's routers, |
| 75 | extract all routes, and return the task name and routes.""" |
| 76 | task = request.param |
| 77 | # Enable development mode to register all routes (including dev-only routes). |
| 78 | monkeypatch.setenv("VLLM_SERVER_DEV_MODE", "1") |
| 79 | |
| 80 | app = FastAPI() |
| 81 | args = Namespace() |
| 82 | app.state = Namespace() |
| 83 | app.state.args = args |
| 84 | |
| 85 | # Register routers for this specific task (development mode already enabled). |
| 86 | register_api_routers( |
| 87 | args, |
| 88 | app, |
| 89 | supported_tasks=(task,), |
| 90 | model_config=MockModelConfig(), |
| 91 | ) |
| 92 | |
| 93 | routes = get_all_http_routes(app) |
| 94 | return task, routes |
| 95 | |
| 96 | |
| 97 | # --------------------------------------------------------------------------- |
| 98 | # Tests for auto-discovered routes |
| 99 | # --------------------------------------------------------------------------- |
| 100 | |
| 101 | |
| 102 | def test_auto_discovered_protected_routes_require_auth(task_routes): |
| 103 | """For every auto-discovered route that starts with a guarded prefix, |
| 104 | verify that authentication is enforced.""" |
| 105 | task, routes = task_routes |
| 106 | app = _create_app_with_mock_routes(routes) |
| 107 | client = TestClient(app) |
| 108 | |
| 109 | for path_template, methods in routes: |
| 110 | if not path_template.startswith(GUARDED_PREFIX): |
| 111 | continue |
| 112 | |
| 113 | test_path = generate_test_path(path_template) |
| 114 | test_method = methods[0] if methods else "GET" |
| 115 | |
| 116 | resp = client.request(test_method, test_path) |
| 117 | assert resp.status_code == 401, ( |
| 118 | f"[{task}] {test_method} {test_path} should reject missing token" |
| 119 | ) |
| 120 | |
| 121 | resp = client.request( |
| 122 | test_method, test_path, headers={"Authorization": "Bearer wrong"} |
| 123 | ) |
| 124 | assert resp.status_code == 401, ( |
| 125 | f"[{task}] {test_method} {test_path} should reject invalid token" |
| 126 | ) |
| 127 | |
| 128 | resp = client.request( |
| 129 | test_method, test_path, headers={"Authorization": "Bearer valid-token"} |
| 130 | ) |
| 131 | assert resp.status_code == 200, ( |
| 132 | f"[{task}] {test_method} {test_path} should accept valid token" |
| 133 | ) |
| 134 | |
| 135 | |
| 136 | def test_auto_discovered_unprotected_routes_no_auth(task_routes): |
| 137 | """For every auto-discovered route that does NOT start with a guarded |
| 138 | prefix, verify that no authentication is required.""" |
| 139 | task, routes = task_routes |
| 140 | app = _create_app_with_mock_routes(routes) |
| 141 | client = TestClient(app) |
| 142 | |
| 143 | for path_template, methods in routes: |
| 144 | if path_template.startswith(GUARDED_PREFIX): |
| 145 | continue |
| 146 | |
| 147 | test_path = generate_test_path(path_template) |
| 148 | test_method = methods[0] if methods else "GET" |
| 149 | |
| 150 | resp = client.request(test_method, test_path) |
| 151 | assert resp.status_code == 200, ( |
| 152 | f"[{task}] {test_method} {test_path} should be accessible without token" |
| 153 | ) |