Coverage for tests/test_run_httpserver.py: 100%
229 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1#!/usr/bin/env python3
3import sys
4import pytest
5import os
6import tempfile
8from http import HTTPStatus
10from unittest.mock import patch
11from unittest.mock import MagicMock
12from unittest.mock import AsyncMock
14from tests.test_utils import temp_sys_path
15from tests.torch_mock import TorchMock
17TMP_DIR = "tmp"
19mock_torch = TorchMock()
22with temp_sys_path("wrapper"):
23 from tests.test_wrapper_model import MockModelGeneration
25with patch.dict(sys.modules, {
26 'nvidia_smi': MagicMock(),
27 'torch': mock_torch,
28 'torch.distributed': MagicMock(),
29 'imageio': MagicMock(),
30 'cv2': MagicMock(),
31 'xfuser': MagicMock(),
32}):
33 with temp_sys_path("wrapper"):
34 import run_httpserver
37@pytest.fixture(autouse=True)
38def clear_models() -> None:
39 run_httpserver.models = {}
42@pytest.mark.asyncio
43async def test_files() -> None:
44 """Check the files endpoints."""
45 app = run_httpserver.app
46 client = app.test_client()
47 response = await client.get("/files")
48 assert response is not None
49 response_json = await response.json
50 assert isinstance(response_json, dict)
52 response = await client.get("/file/filename.txt")
53 assert response is not None
54 assert response.status_code == HTTPStatus.NOT_FOUND
55 response_json = await response.json
56 assert isinstance(response_json, dict)
57 assert "error" in response_json
58 assert response_json["error"] == "File not found"
60 # Create temp file in /tmp
61 with tempfile.NamedTemporaryFile(mode="w", delete=False, dir="/tmp", suffix=".txt") as tmp_file:
62 tmp_file.write("Test file content")
63 tmp_filename = tmp_file.name
64 filename = os.path.basename(tmp_filename)
66 # Download
67 response = await client.get(f"/file/{filename}")
68 assert response is not None
69 assert response.status_code == HTTPStatus.OK
70 content = await response.get_data(as_text=True)
71 assert content == "Test file content"
73 # Info
74 response = await client.get(f"/file_info/{filename}")
75 assert response is not None
76 assert response.status_code == HTTPStatus.OK
77 response_json = await response.json
78 assert len(response_json) > 0
79 assert response_json["size"] == 17
80 assert response_json["mimetype"] == "text/plain"
81 assert response_json["name"].endswith(".txt")
82 assert response_json["date"] > 0
83 assert response_json["type"] == "text"
84 # get_text_file_info()
85 assert response_json["num_chars"] == 17
86 assert response_json["num_words"] == 3
87 assert response_json["num_lines"] == 1
89 os.unlink(tmp_filename)
92@pytest.mark.timeout(5)
93@pytest.mark.asyncio
94async def test_yolo() -> None:
95 """Check the YOLO endpoint."""
96 app = run_httpserver.app
97 client = app.test_client()
98 response = await client.post('/yolo')
99 assert response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
102@pytest.mark.timeout(5)
103@pytest.mark.asyncio
104async def test_transcript() -> None:
105 """Check the transcript endpoint."""
106 app = run_httpserver.app
107 client = app.test_client()
108 response = await client.post("/podcasttranscript")
109 assert response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
110 response_json = await response.json
111 assert isinstance(response_json, dict)
112 assert "error" in response_json
113 assert response_json["error"] == "Podcast transcript model not initialized"
115 run_httpserver.models = {"podcasttranscript": MockModelGeneration()}
116 response = await client.post("/podcasttranscript")
117 assert response.status_code == HTTPStatus.BAD_REQUEST
118 response_json = await response.json
119 assert isinstance(response_json, dict)
120 assert "error" in response_json
121 assert response_json["error"] == "No JSON body received"
123 # Setting up the mock model
124 run_httpserver.models = {"podcasttranscript": MockModelGeneration()}
126 # Transcript mocked
127 response = await client.post(
128 "/podcasttranscript",
129 json={"args": {}})
130 assert response.status_code == HTTPStatus.BAD_REQUEST
131 response_json = await response.json
132 assert isinstance(response_json, dict)
133 assert "error" in response_json
134 assert response_json["error"] == "Model not initialized. Current status: initializing."
136 # Transcript stream
137 response = await client.post("/podcasttranscript/stream")
138 assert response.status_code == HTTPStatus.BAD_REQUEST
140 with pytest.raises(AssertionError):
141 response = await client.post(
142 "/podcasttranscript/stream",
143 json={
144 "args": {
145 "prompt": "Test prompt"
146 }
147 })
148 assert response.status_code == HTTPStatus.BAD_REQUEST
151@pytest.mark.asyncio
152async def test_kokoro() -> None:
153 """Check the kokoro endpoint."""
154 app = run_httpserver.app
155 client = app.test_client()
156 response = await client.post("/kokoro")
157 assert response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
160@pytest.mark.asyncio
161async def test_service_endpoints() -> None:
162 app = run_httpserver.app
163 client = app.test_client()
164 service_names = [
165 "flux", "fluxkontext", "fluxkrea", "fluxupscaler", "fluxupscaler/video",
166 "fantasytalking", "hunyuanavatar",
167 "kokoro", "xtts", "vibevoice",
168 "hidream",
169 "realesrgan", "realesrgan/video",
170 "qwenimage", "qwenimageedit",
171 "realesrgan", "imageresize",
172 "bagel", "llamagen", "januspro",
173 "thinksound", "dia",
174 "ltx", "wan", "wan22",
175 "hunyuanframepack", "hunyuanframepackf1",
176 ]
177 for service_name in service_names:
178 response = await client.post(f"/{service_name}")
179 status = response.status_code
180 assert status == HTTPStatus.INTERNAL_SERVER_ERROR, f"Service {service_name} failed with {status}"
183@pytest.mark.asyncio
184async def test_service_health() -> None:
185 """Check the service health endpoint."""
186 app = run_httpserver.app
187 client = app.test_client()
188 response = await client.get("/wan/health")
189 assert response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
192@pytest.mark.asyncio
193async def test_mock_model() -> None:
194 """Check the service health endpoint."""
195 app = run_httpserver.app
196 client = app.test_client()
198 # Setup mock model
199 run_httpserver.models = {"mockmodel": MockModelGeneration()}
201 # Health
202 response = await client.get("/mockmodel/health")
203 assert response.status_code == HTTPStatus.OK
204 model_health = await response.json
205 assert len(model_health) > 0
206 assert "gen_timer" in model_health
209@pytest.mark.asyncio
210async def test_yolo_mock() -> None:
211 """Check YOLO using mocks."""
212 app = run_httpserver.app
213 client = app.test_client()
215 # Setup mock model for YOLO
216 run_httpserver.models = {
217 "yolo": MockModelGeneration(output_type="list_pillow")
218 }
220 # Health
221 response = await client.get("/yolo/health")
222 assert response.status_code == HTTPStatus.OK
223 model_health = await response.json
224 assert len(model_health) > 0
225 assert "gen_timer" in model_health
227 # Generate with no body
228 response = await client.post("/yolo")
229 assert response.status_code == HTTPStatus.BAD_REQUEST
230 response_json = await response.json
231 assert "error" in response_json
232 error_msg = response_json["error"]
233 assert "No JSON body received" in error_msg
235 # Generate with wrong body
236 response = await client.post(
237 "/yolo",
238 json={"foo": "bar"})
239 assert response.status_code == HTTPStatus.BAD_REQUEST
241 # Generate without arguments
242 response = await client.post(
243 "/yolo",
244 json={
245 "args": {},
246 "job_id": "job0",
247 },
248 headers={"Content-Type": "application/json"})
249 assert response.status_code == HTTPStatus.BAD_REQUEST
250 # TODO need to fix this
253@pytest.mark.asyncio
254async def test_index() -> None:
255 """Check the index endpoint."""
256 app = run_httpserver.app
257 client = app.test_client()
258 response = await client.get("/")
259 assert response.status_code == HTTPStatus.OK
260 index_html = await response.get_data(as_text=True)
261 assert index_html is not None
262 assert "LMM Models" in index_html
265@pytest.mark.asyncio
266async def test_timestamps() -> None:
267 """Check the timestamps endpoint."""
268 app = run_httpserver.app
269 client = app.test_client()
270 response = await client.get("/timestamps")
271 assert response.status_code == HTTPStatus.OK
272 timestamps = await response.json
273 assert timestamps == {}
275 run_httpserver.models = {"mockmodel": MockModelGeneration()}
276 response = await client.get("/timestamps")
277 assert response.status_code == HTTPStatus.OK
278 timestamps = await response.json
279 assert timestamps is not None
280 assert "mockmodel" in timestamps
283@pytest.mark.asyncio
284async def test_health() -> None:
285 """Check the health endpoint."""
286 app = run_httpserver.app
287 client = app.test_client()
288 response = await client.get("/health")
289 assert response.status_code == HTTPStatus.OK
290 health = await response.json
291 assert health == {}
293 run_httpserver.models = {"mockmodel": MockModelGeneration()}
294 response = await client.get("/health")
295 assert response.status_code == HTTPStatus.OK
296 health = await response.json
297 assert health is not None
298 assert "mockmodel" in health
301def test_setup_dist_environment_non_mig() -> None:
302 """setup_dist_environment uses local_rank as device_id when multiple GPUs are visible (non-MIG)."""
303 env = {
304 "MASTER_ADDR": "localhost",
305 "MASTER_PORT": "12355",
306 "RANK": "1",
307 "LOCAL_RANK": "1",
308 "WORLD_SIZE": "2",
309 "LOCAL_WORLD_SIZE": "2",
310 }
311 with patch.dict(os.environ, env, clear=False):
312 mock_torch.cuda.is_available.return_value = True
313 mock_torch.cuda.device_count.return_value = 2
314 run_httpserver.setup_dist_environment()
315 mock_torch.cuda.set_device.assert_called_with(1)
318def test_setup_dist_environment_mig() -> None:
319 """setup_dist_environment uses device 0 when MIG restricts container to a single visible device."""
320 env = {
321 "MASTER_ADDR": "localhost",
322 "MASTER_PORT": "12355",
323 "RANK": "1",
324 "LOCAL_RANK": "1",
325 "WORLD_SIZE": "2",
326 "LOCAL_WORLD_SIZE": "2",
327 }
328 with patch.dict(os.environ, env, clear=False):
329 mock_torch.cuda.is_available.return_value = True
330 # MIG: each process sees only its own MIG instance (device_count=1)
331 mock_torch.cuda.device_count.return_value = 1
332 run_httpserver.setup_dist_environment()
333 # local_rank=1 is out of range; should fall back to device 0
334 mock_torch.cuda.set_device.assert_called_with(0)
337@pytest.mark.asyncio
338@pytest.mark.timeout(10)
339async def test_run_httpserver_http(monkeypatch: pytest.MonkeyPatch) -> None:
340 """run_httpserver() configures Hypercorn without SSL when no certfile is given."""
341 mock_serve = AsyncMock()
342 monkeypatch.setattr(run_httpserver, "serve", mock_serve)
343 await run_httpserver.run_httpserver(host="127.0.0.1", port=9999)
344 mock_serve.assert_awaited_once()
345 config = mock_serve.call_args[0][1]
346 assert config.bind == ["127.0.0.1:9999"]
347 # No SSL attributes set
348 assert getattr(config, "certfile", None) is None
349 assert getattr(config, "keyfile", None) is None
352@pytest.mark.asyncio
353@pytest.mark.timeout(10)
354async def test_run_httpserver_https(monkeypatch: pytest.MonkeyPatch) -> None:
355 """run_httpserver() configures Hypercorn with SSL when certfile and keyfile are given."""
356 mock_serve = AsyncMock()
357 monkeypatch.setattr(run_httpserver, "serve", mock_serve)
358 await run_httpserver.run_httpserver(
359 host="127.0.0.1",
360 port=9999,
361 certfile="/tmp/cert.pem",
362 keyfile="/tmp/key.pem",
363 )
364 mock_serve.assert_awaited_once()
365 config = mock_serve.call_args[0][1]
366 assert config.bind == ["127.0.0.1:9999"]
367 assert config.certfile == "/tmp/cert.pem"
368 assert config.keyfile == "/tmp/key.pem"
371def test_arg_parsing_https_args() -> None:
372 """arg_parsing() accepts --certfile and --keyfile arguments."""
373 with patch("sys.argv", ["run_httpserver", "--mock", "--certfile", "/tmp/cert.pem", "--keyfile", "/tmp/key.pem"]):
374 args, _ = run_httpserver.arg_parsing()
375 assert args.certfile == "/tmp/cert.pem"
376 assert args.keyfile == "/tmp/key.pem"
379def test_arg_parsing_https_defaults() -> None:
380 """arg_parsing() defaults certfile and keyfile to None (HTTP mode)."""
381 with patch("sys.argv", ["run_httpserver", "--mock"]):
382 args, _ = run_httpserver.arg_parsing()
383 assert args.certfile is None
384 assert args.keyfile is None