Coverage for tests/test_run_httpserver_mock.py: 99%
255 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 asyncio
6import logging
8from pytest import raises
10from http import HTTPStatus
12from unittest.mock import patch
13from unittest.mock import MagicMock
14from unittest.mock import AsyncMock
15from unittest.mock import ANY
17from tests.torch_mock import TorchMock
19from streamwise_apps import STREAMWISE_APPS
20from streamwise_apps import VLLM_SERVICES
22mock_torch = TorchMock()
24sys.path.append("wrapper")
26# Quart mocks
27mock_quart = MagicMock()
28mock_quart.route = lambda *args, **kwargs: (lambda f: f)
29mock_quart.render_template = AsyncMock(return_value="mocked_template")
30mock_quart.send_file = AsyncMock(return_value="mocked_file")
31mock_quart.send_from_directory = AsyncMock(return_value="mocked_file")
32mock_quart.jsonify = lambda x: x
33mock_app = MagicMock()
34mock_quart.Quart = MagicMock(return_value=mock_app)
35mock_quart.Response = lambda *args, **kwargs: {"response": args, "kwargs": kwargs}
36mock_request = MagicMock()
37mock_request.args = {}
38mock_request.json = {}
39mock_request.get_json = AsyncMock(return_value={})
40mock_quart.request = mock_request
42mock_serve = AsyncMock()
44mock_modules = {
45 'nvidia_smi': MagicMock(),
46 'torch': mock_torch,
47 'hypercorn': MagicMock(),
48 'hypercorn.config': MagicMock(),
49 'hypercorn.asyncio': MagicMock(serve=mock_serve),
50 'quart': mock_quart,
51 "quart.request": mock_quart.request,
52 "quart.jsonify": mock_quart.jsonify,
53 "quart.send_file": mock_quart.send_file,
54 "quart.send_from_directory": mock_quart.send_from_directory,
55 "quart.render_template": mock_quart.render_template,
56 "quart.route": mock_quart.route,
57 "quart.Response": mock_quart.Response,
58 'imageio': MagicMock(),
59 'cv2': MagicMock(),
60 'xfuser': MagicMock(),
61 'sklearn': MagicMock(),
62 'scipy': MagicMock(),
63 'scipy.stats': MagicMock(),
64}
65mock_modules.update(mock_torch.get_sub_modules())
67with patch.dict(sys.modules, mock_modules):
68 from run_httpserver import main
69 from run_httpserver import arg_parsing
70 from run_httpserver import send_task
71 from run_httpserver import nccl_worker
72 from run_httpserver import wait_for_everybody
73 from run_httpserver import index
74 from run_httpserver import gen_img
75 from run_httpserver import gen_video
76 from run_httpserver import gen_audio
77 from run_httpserver import get_service_names
78 from run_httpserver import setup_dist_environment
79 import run_httpserver as _run_httpserver
82def test_get_service_names() -> None:
83 """Check that we have the expected services covered."""
84 service_names = get_service_names()
85 assert len(service_names) == 48
86 assert "mock" in service_names
87 assert "wan" in service_names
88 assert "wan22" in service_names
89 assert "fantasytalking" in service_names
90 assert "flux" in service_names
91 assert "fluxkrea" in service_names
92 assert "qwenimage" in service_names
93 assert "qwenimageedit" in service_names
94 assert "hunyuanimage" in service_names
95 assert "yolo" in service_names
96 assert "kokoro" in service_names
97 assert "vibevoice" in service_names
98 assert "notexisting" not in service_names
100 for vllm_service in VLLM_SERVICES:
101 if vllm_service != "llm":
102 assert vllm_service in service_names
104 # Applications
105 for streamwise_app in STREAMWISE_APPS:
106 assert streamwise_app in service_names
109@pytest.mark.asyncio
110async def test_server_help() -> None:
111 """Check the HTTP server start."""
112 test_args = ["run_http_server.py", "--help"]
113 with patch.object(sys, "argv", test_args):
114 with raises(SystemExit) as exc_info:
115 await main()
116 assert exc_info.value.code == 0
119# This test is expensive; we swallow the timeout
120# @pytest.mark.timeout(2)
121@pytest.mark.asyncio
122async def test_server() -> None:
123 """Check the base HTTP server start."""
124 test_args = ["run_http_server.py"]
125 with patch.object(sys, "argv", test_args):
126 try:
127 # Use asyncio.wait_for to limit the runtime to 2s
128 await asyncio.wait_for(main(), timeout=2)
129 except (asyncio.TimeoutError, asyncio.CancelledError):
130 print("HTTP server finished (swallowed)")
133@pytest.mark.asyncio
134async def test_server_mock() -> None:
135 """Check the HTTP server start for the mock service."""
136 test_args = ["run_http_server.py", "--mock"]
137 with patch.object(sys, "argv", test_args):
138 await main()
141@pytest.mark.asyncio
142async def test_server_imageresize() -> None:
143 """Check the HTTP server start for the image resize service."""
144 test_args = ["run_http_server.py", "--imageresize"]
145 with patch.object(sys, "argv", test_args):
146 await main()
149@pytest.mark.asyncio
150@pytest.mark.parametrize("model_name", [
151 "flux",
152 "hidream",
153 "fantasytalking",
154 "hunyuanframepack",
155 "hunyuanframepackf1",
156 "wan",
157 "wan21",
158 "wan22",
159])
160async def test_server_service_diffusers(model_name: str) -> None:
161 """Check the HTTP server start for multiple services.
162 These services fail with:
163 RuntimeError: Failed to import diffusers.pipelines.pipeline_utils
164 """
165 test_args = ["run_http_server.py", f"--{model_name}"]
166 with patch.object(sys, "argv", test_args):
167 with raises(
168 (ModuleNotFoundError, RuntimeError),
169 match="(diffusers|This module requires CUDA support|Found no NVIDIA driver)"
170 ):
171 await main()
174@pytest.mark.asyncio
175@pytest.mark.parametrize("model_name", [
176 "realesrgan",
177 "yolo", # ultralytics
178 "kokoro", # KPipeline
179 "xtts", # TTS
180 # "vibevoice", # demo
181 "podcasttranscript", # fitz
182 # "januspro" # janus
183])
184async def test_server_service_import(model_name: str) -> None:
185 """Check the HTTP server start for multiple services.
186 They fail with the following errors:
187 ModuleNotFoundError: No module named 'RealESRGAN'
188 ModuleNotFoundError: No module named 'fitz'
189 ModuleNotFoundError: No module named 'janus'
190 ImportError: cannot import name 'KPipeline' from 'kokoro'
191 """
192 test_args = ["run_http_server.py", f"--{model_name}"]
193 with patch.object(sys, "argv", test_args):
194 with pytest.raises((ModuleNotFoundError, ImportError)):
195 await main()
198@pytest.mark.asyncio
199async def test_server_all_services() -> None:
200 """Check the HTTP server start for all services."""
201 service_names = get_service_names()
202 for vllm_service in VLLM_SERVICES:
203 if vllm_service in service_names:
204 service_names.remove(vllm_service) # Separate vLLM service
205 for streamwise_app in STREAMWISE_APPS:
206 if streamwise_app in service_names:
207 service_names.remove(streamwise_app) # Separate service
208 service_names.remove("streamwise") # Separate service
210 for model_name in service_names:
211 try:
212 test_args = ["run_http_server.py", f"--{model_name}"]
213 with patch.object(sys, "argv", test_args):
214 try:
215 logging.info(f"Testing service: {model_name}")
216 # Use asyncio.wait_for to limit the runtime to 2s
217 await asyncio.wait_for(main(), timeout=2.0)
218 except (asyncio.TimeoutError, asyncio.CancelledError):
219 logging.error(f"HTTP server for {model_name} finished (swallowed)")
220 except (
221 ModuleNotFoundError,
222 ImportError,
223 RuntimeError,
224 TypeError, # Janus Pro
225 ):
226 logging.warning(f"Skipping service {model_name} due to error")
227 finally:
228 loop = asyncio.get_running_loop()
229 executor = getattr(loop, "_default_executor", None)
230 if executor:
231 executor.shutdown(wait=False, cancel_futures=True)
234@pytest.mark.asyncio
235async def test_server_wrong_server() -> None:
236 """Check the HTTP server start with a wrong argument."""
237 with raises(SystemExit) as exc_info:
238 test_args = ["run_http_server.py", "--wrong"]
239 with patch.object(sys, "argv", test_args):
240 await main()
241 assert exc_info.value.code == 2
244def test_arg_parsing() -> None:
245 """Check the argument parsing."""
246 test_args = ["run_http_server.py", "--wan", "--flux", "--hunyuanframepackf1"]
247 with patch.object(sys, "argv", test_args):
248 args, engine_config = arg_parsing()
249 assert args.wan is True
250 assert args.flux is True
251 assert args.fluxupscaler is False
252 assert args.fluxkontext is False
253 assert args.fluxkrea is False
254 assert args.hidream is False
255 assert args.qwenimage is False
256 assert args.qwenimageedit is False
257 assert args.realesrgan is False
258 assert args.kokoro is False
259 assert args.hunyuanframepack is False
260 assert args.hunyuanframepackf1 is True
261 assert args.hunyuanframepackvae is False
262 assert engine_config is None
264 test_args = ["run_http_server.py", "--nonexisting"]
265 with patch.object(sys, "argv", test_args):
266 with raises(SystemExit) as exc_info:
267 args, engine_config = arg_parsing()
268 assert exc_info.value.code == 2
271@pytest.mark.asyncio
272async def test_send_task() -> None:
273 """Basic call to cover the function."""
274 await send_task({})
277@pytest.mark.asyncio
278async def test_nccl_worker() -> None:
279 """Basic call to cover the function."""
280 await nccl_worker()
283@pytest.mark.asyncio
284async def test_wait_for_everybody() -> None:
285 """Basic call to cover the function."""
286 await wait_for_everybody()
289@pytest.mark.asyncio
290async def test_content() -> None:
291 """Check the HTTP server content."""
292 with raises(TypeError):
293 # TODO fix: TypeError: object MagicMock can't be used in 'await' expression
294 assert await index() == "index page"
295 # assert await health() == {}
296 # assert await model_health("dummy") == {"ok": True}
299@pytest.mark.asyncio
300async def test_gen_img() -> None:
301 """Check the image generation endpoint."""
302 response = await gen_img(None)
303 assert response == ({"error": "Not initialized"}, HTTPStatus.INTERNAL_SERVER_ERROR)
305 mock_model = AsyncMock()
306 mock_model.status = "ok"
307 mock_model.running = True
308 response = await gen_img(mock_model)
309 assert response == ({"error": "Generation in progress"}, HTTPStatus.SERVICE_UNAVAILABLE)
311 mock_model.running = False
312 response = await gen_img(mock_model)
313 assert response is not None
314 # assert response == "mocked_file"
317@pytest.mark.asyncio
318async def test_gen_video() -> None:
319 """Check the video generation endpoint."""
320 response = await gen_video(None)
321 assert response == ({"error": "Not initialized"}, HTTPStatus.INTERNAL_SERVER_ERROR)
323 mock_model = MagicMock()
324 mock_model.status = "ok"
325 mock_model.generate = AsyncMock(return_value="mocked_file.mp4")
326 mock_model.get_rest_args = AsyncMock(return_value={
327 "task": "mock",
328 "args": {}
329 })
330 mock_model.running = True
331 response = await gen_video(mock_model)
332 assert response == ({"error": "Generation in progress"}, HTTPStatus.SERVICE_UNAVAILABLE)
334 mock_model.running = False
335 response = await gen_video(mock_model)
336 assert response is not None
337 response_msg, code = response
338 assert code == HTTPStatus.INTERNAL_SERVER_ERROR
339 assert response_msg["error"] == "Video file not found: mocked_file.mp4"
342@pytest.mark.asyncio
343async def test_gen_audio() -> None:
344 """Check the audio generation endpoint."""
345 response = await gen_audio(None)
346 assert response == ({"error": "Not initialized"}, HTTPStatus.INTERNAL_SERVER_ERROR)
348 # Blocked generation
349 mock_model = MagicMock()
350 mock_model.status = "ok"
351 mock_model.generate = AsyncMock(return_value="mocked_file.wav")
352 mock_model.get_rest_args = AsyncMock(return_value={
353 "task": "mock",
354 "args": {}
355 })
356 mock_model.running = True
357 response = await gen_audio(mock_model)
358 assert response == ({"error": "Generation in progress"}, HTTPStatus.SERVICE_UNAVAILABLE)
360 # Missing file
361 mock_model.running = False
362 response = await gen_audio(mock_model)
363 assert response is not None
364 response_msg, code = response
365 assert code == HTTPStatus.INTERNAL_SERVER_ERROR
366 assert response_msg["error"] == "Audio file not found: mocked_file.wav"
368 # No output
369 mock_model.generate = AsyncMock(return_value=None)
370 response = await gen_audio(mock_model)
371 assert response is not None
372 response_msg, code = response
373 assert code == HTTPStatus.INTERNAL_SERVER_ERROR
374 assert response_msg["error"] == "No audio generated"
376 # Use existing audio file
377 mock_quart.send_from_directory.reset_mock()
378 mock_model.generate = AsyncMock(return_value="tests/data/audio_4675.wav")
379 response = await gen_audio(mock_model)
380 assert response is not None
381 assert response == "mocked_file" # send_from_directory() returns this
382 mock_quart.send_from_directory.assert_awaited_once_with(
383 "tests/data",
384 "audio_4675.wav",
385 mimetype="audio/wav",
386 as_attachment=True,
387 attachment_filename=ANY,
388 )
389 mock_quart.send_file.assert_not_awaited()
392def test_setup_dist_environment_mig_warning(caplog: pytest.LogCaptureFixture) -> None:
393 """
394 When world_size > visible CUDA devices (MIG partition case), setup_dist_environment
395 must log a warning and clamp world_size to the number of visible devices.
396 The TorchMock returns device_count=1, so setting WORLD_SIZE=2 triggers the path.
397 """
398 import os
400 saved_env = {
401 k: os.environ.get(k)
402 for k in ("MASTER_ADDR", "MASTER_PORT", "RANK", "LOCAL_RANK",
403 "NODE_RANK", "WORLD_SIZE", "LOCAL_WORLD_SIZE", "NPROC_PER_NODE")
404 }
405 try:
406 os.environ["MASTER_ADDR"] = "localhost"
407 os.environ["MASTER_PORT"] = "12355"
408 os.environ["RANK"] = "1"
409 os.environ["LOCAL_RANK"] = "1"
410 os.environ["NODE_RANK"] = "0"
411 os.environ["WORLD_SIZE"] = "2" # 2 processes, but only 1 visible device
412 os.environ["LOCAL_WORLD_SIZE"] = "2"
414 with caplog.at_level(logging.WARNING):
415 setup_dist_environment()
417 # world_size must be clamped down to the number of visible CUDA devices (1)
418 assert _run_httpserver.world_size == 1
419 assert _run_httpserver.local_world_size == 1
420 # And a MIG-related warning must have been logged
421 assert any(
422 "world_size=2" in record.message and "MIG" in record.message
423 for record in caplog.records
424 ), f"Expected MIG warning, got: {[r.message for r in caplog.records]}"
425 finally:
426 # Restore original env and module globals
427 for k, v in saved_env.items():
428 if v is None:
429 os.environ.pop(k, None)
430 else:
431 os.environ[k] = v
432 _run_httpserver.rank = 0
433 _run_httpserver.world_size = 1
434 _run_httpserver.local_world_size = 1