@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 161 lines 4.53 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for middleware that's off by default and can be toggled through
4 server arguments, mainly --api-key and --enable-request-id-headers.
5 """
6
7 from http import HTTPStatus
8
9 import pytest
10 import requests
11
12 from tests.utils import RemoteOpenAIServer
13
14 # Use a small embeddings model for faster startup and smaller memory footprint.
15 # Since we are not testing any chat functionality,
16 # using a chat capable model is overkill.
17 MODEL_NAME = "intfloat/multilingual-e5-small"
18
19
20 @pytest.fixture(scope="module")
21 def server(request: pytest.FixtureRequest):
22 passed_params = []
23 if hasattr(request, "param"):
24 passed_params = request.param
25 if isinstance(passed_params, str):
26 passed_params = [passed_params]
27
28 args = [
29 "--runner",
30 "pooling",
31 # use half precision for speed and memory savings in CI environment
32 "--dtype",
33 "float16",
34 "--max-model-len",
35 "512",
36 "--enforce-eager",
37 "--max-num-seqs",
38 "2",
39 *passed_params,
40 ]
41
42 with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
43 yield remote_server
44
45
46 @pytest.mark.asyncio
47 async def test_no_api_token(server: RemoteOpenAIServer):
48 response = requests.get(server.url_for("v1/models"))
49 assert response.status_code == HTTPStatus.OK
50
51
52 @pytest.mark.asyncio
53 async def test_no_request_id_header(server: RemoteOpenAIServer):
54 response = requests.get(server.url_for("health"))
55 assert "X-Request-Id" not in response.headers
56
57
58 @pytest.mark.parametrize(
59 "server",
60 [["--api-key", "test"]],
61 indirect=True,
62 )
63 @pytest.mark.asyncio
64 async def test_missing_api_token(server: RemoteOpenAIServer):
65 response = requests.get(server.url_for("v1/models"))
66 assert response.status_code == HTTPStatus.UNAUTHORIZED
67
68
69 @pytest.mark.parametrize(
70 "server",
71 [["--api-key", "test"]],
72 indirect=True,
73 )
74 @pytest.mark.asyncio
75 async def test_passed_api_token(server: RemoteOpenAIServer):
76 response = requests.get(
77 server.url_for("v1/models"), headers={"Authorization": "Bearer test"}
78 )
79 assert response.status_code == HTTPStatus.OK
80
81
82 @pytest.mark.parametrize(
83 "server",
84 [["--api-key", "test"]],
85 indirect=True,
86 )
87 @pytest.mark.asyncio
88 async def test_not_v1_or_v2_path_skips_auth(server: RemoteOpenAIServer):
89 # Authorization check is skipped for paths that
90 # don't start with /v1 or /v2 (e.g. /health, /metrics).
91 response = requests.get(server.url_for("health"))
92 assert response.status_code == HTTPStatus.OK
93
94
95 # ---------------------------------------------------------------------------
96 # /v2 path authentication tests
97 # ---------------------------------------------------------------------------
98
99
100 @pytest.mark.parametrize(
101 "server",
102 [["--api-key", "test"]],
103 indirect=True,
104 )
105 @pytest.mark.asyncio
106 async def test_v2_endpoint_rejects_missing_api_token(server: RemoteOpenAIServer):
107 # /v2/embed should require authentication when --api-key is set.
108 body = {
109 "model": MODEL_NAME,
110 "texts": ["hello"],
111 "embedding_types": ["float"],
112 }
113 response = requests.post(server.url_for("/v2/embed"), json=body)
114 assert response.status_code == HTTPStatus.UNAUTHORIZED
115
116
117 @pytest.mark.parametrize(
118 "server",
119 [["--api-key", "test"]],
120 indirect=True,
121 )
122 @pytest.mark.asyncio
123 async def test_v2_endpoint_accepts_valid_api_token(server: RemoteOpenAIServer):
124 # /v2/embed should accept requests with a valid API key.
125 body = {
126 "model": MODEL_NAME,
127 "texts": ["hello"],
128 "embedding_types": ["float"],
129 }
130 response = requests.post(
131 server.url_for("/v2/embed"),
132 json=body,
133 headers={"Authorization": "Bearer test"},
134 )
135 assert response.status_code == HTTPStatus.OK
136
137
138 @pytest.mark.parametrize(
139 "server",
140 ["--enable-request-id-headers"],
141 indirect=True,
142 )
143 @pytest.mark.asyncio
144 async def test_enable_request_id_header(server: RemoteOpenAIServer):
145 response = requests.get(server.url_for("health"))
146 assert "X-Request-Id" in response.headers
147 assert len(response.headers.get("X-Request-Id", "")) == 32
148
149
150 @pytest.mark.parametrize(
151 "server",
152 ["--enable-request-id-headers"],
153 indirect=True,
154 )
155 @pytest.mark.asyncio
156 async def test_custom_request_id_header(server: RemoteOpenAIServer):
157 response = requests.get(
158 server.url_for("health"), headers={"X-Request-Id": "Custom"}
159 )
160 assert "X-Request-Id" in response.headers
161 assert response.headers.get("X-Request-Id") == "Custom"