From e0cdbe9a5d7ec654bbbe26c2fb6e76abe41446d2 Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Mon, 13 Apr 2026 19:21:42 +0800 Subject: [PATCH] [Misc] Cleanup: use consistent pytest-mock in unit tests (#2698) Signed-off-by: yuanheng --- tests/comfyui/conftest.py | 18 +- tests/comfyui/test_comfyui_integration.py | 93 +- .../test_generation_scheduler_restore.py | 27 +- .../test_distributed_vae_executor.py | 41 +- .../models/bagel/test_trajectory_recording.py | 34 +- .../models/flux2/test_flux2_transformer_tp.py | 20 +- .../offloader/test_sequential_backend.py | 114 +- .../quantization/test_int8_config.py | 32 +- tests/diffusion/test_diffusion_scheduler.py | 103 +- .../diffusion/test_diffusion_step_pipeline.py | 26 +- .../test_diffusion_worker_cuda_profiler.py | 6 +- .../test_multiproc_engine_concurrency.py | 28 +- tests/engine/test_arg_utils.py | 9 +- tests/engine/test_async_omni_engine_input.py | 15 +- .../engine/test_async_omni_engine_outputs.py | 20 +- tests/engine/test_single_stage_mode.py | 1527 ++++++++++------- .../openai_api/test_serving_chat_speaker.py | 40 +- .../openai_api/test_serving_speech.py | 211 ++- .../openai_api/test_serving_speech_stream.py | 117 +- tests/entrypoints/test_omni_base_profiler.py | 25 +- tests/entrypoints/test_serve.py | 184 +- .../test_mimo_audio_code2wav_batch_decode.py | 40 +- .../qwen2_5_omni/test_qwen2_5_omni_embed.py | 37 +- .../qwen3_tts/test_code_predictor_dtype.py | 131 +- .../models/test_fish_speech_voice_cache.py | 30 +- tests/test_fish_speech_voice_cache.py | 39 +- 26 files changed, 1598 insertions(+), 1369 deletions(-) diff --git a/tests/comfyui/conftest.py b/tests/comfyui/conftest.py index 0b4565e94..4280d3506 100644 --- a/tests/comfyui/conftest.py +++ b/tests/comfyui/conftest.py @@ -9,8 +9,8 @@ This module sets up the test environment by: import os import sys +from types import ModuleType, SimpleNamespace from typing import BinaryIO, TypedDict -from unittest.mock import MagicMock def pytest_configure(config): @@ -58,15 +58,15 @@ def _setup_comfyui_test_environment(): else: file.write(self._data) - mock_comfy_api = MagicMock() - mock_comfy_api_input = MagicMock() + mock_comfy_api = ModuleType("comfy_api") + mock_comfy_api_input = ModuleType("comfy_api.input") mock_comfy_api_input.AudioInput = AudioInput mock_comfy_api_input.VideoInput = VideoInput mock_comfy_api.input = mock_comfy_api_input - mock_comfy_api_latest = MagicMock() - mock_comfy_api_latest.Types.VideoComponents = MagicMock(side_effect=lambda **kwargs: kwargs) - mock_comfy_api_latest.InputImpl.VideoFromComponents = MagicMock( - side_effect=lambda _: VideoInput(b"mock_video_from_components") + mock_comfy_api_latest = ModuleType("comfy_api.latest") + mock_comfy_api_latest.Types = SimpleNamespace(VideoComponents=lambda **kwargs: kwargs) + mock_comfy_api_latest.InputImpl = SimpleNamespace( + VideoFromComponents=lambda _: VideoInput(b"mock_video_from_components") ) mock_comfy_api.latest = mock_comfy_api_latest @@ -76,8 +76,8 @@ def _setup_comfyui_test_environment(): sample_rate = 24000 return waveform, sample_rate - mock_comfy_extras = MagicMock() - mock_nodes_audio = MagicMock() + mock_comfy_extras = ModuleType("comfy_extras") + mock_nodes_audio = ModuleType("comfy_extras.nodes_audio") mock_nodes_audio.load = mock_load mock_comfy_extras.nodes_audio = mock_nodes_audio diff --git a/tests/comfyui/test_comfyui_integration.py b/tests/comfyui/test_comfyui_integration.py index f6ce82f9b..80e86d824 100644 --- a/tests/comfyui/test_comfyui_integration.py +++ b/tests/comfyui/test_comfyui_integration.py @@ -13,7 +13,6 @@ from collections.abc import Iterable, Sequence from enum import StrEnum, auto from types import SimpleNamespace from typing import Any, NamedTuple -from unittest.mock import AsyncMock, MagicMock, patch import pytest import requests @@ -28,6 +27,7 @@ from comfyui_vllm_omni.nodes import ( ) from comfyui_vllm_omni.utils.types import AutoregressionSamplingParams, DiffusionSamplingParams, WanModelSpecificParams from PIL import Image +from pytest_mock import MockerFixture from vllm import SamplingParams from vllm.outputs import CompletionOutput, RequestOutput from vllm.utils.argparse_utils import FlexibleArgumentParser @@ -217,9 +217,10 @@ def _build_diffusion_video_output() -> OmniRequestOutput: def _build_diffusion_image_output_for_chat_endpoint() -> OmniRequestOutput: - request_output = MagicMock() - request_output.images = [_build_image_output(color="blue")] - request_output.finished = True + request_output = SimpleNamespace( + images=[_build_image_output(color="blue")], + finished=True, + ) return OmniRequestOutput( request_id="test_req_img_chat", finished=True, @@ -389,51 +390,55 @@ def sampling_case(request) -> SamplingCase: @pytest.fixture -def mock_async_omni(server_case: ServerCase, sampling_case: SamplingCase): +def mock_async_omni( + server_case: ServerCase, + sampling_case: SamplingCase, + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, +): async def _mock_preprocess_chat(self, *args, **kwargs): return ([{"role": "user", "content": "test"}], [{"prompt": "test prompt"}]) # Need to mock AsyncOmni itself (not only its generate method) because # 1. The API layer uses its stage_list and stage_configs attributes # 2. Its __init__ method has slow side effects (model & config loading). - with ( - patch("vllm_omni.entrypoints.openai.api_server.AsyncOmni") as MockAsyncOmni, - patch( - "vllm_omni.entrypoints.openai.serving_chat.OmniOpenAIServingChat._preprocess_chat", - new=_mock_preprocess_chat, - ), - ): - mock_instance = AsyncMock(spec=RealAsyncOmni) - mock_instance.generate = _build_mock_outputs(server_case.outputs, sampling_case, server_case) + mock_async_omni_cls = mocker.patch("vllm_omni.entrypoints.openai.api_server.AsyncOmni") + monkeypatch.setattr( + "vllm_omni.entrypoints.openai.serving_chat.OmniOpenAIServingChat._preprocess_chat", + _mock_preprocess_chat, + ) - mock_instance.stage_list = server_case.stage_list - mock_instance.stage_configs = server_case.stage_configs - mock_instance.output_modalities = _build_output_modalities(server_case.stage_configs) - mock_instance.default_sampling_params_list = [ - SamplingParams() if _stage_type(stage) != "diffusion" else MagicMock() - for stage in server_case.stage_configs - ] - mock_instance.errored = False - mock_instance.dead_error = RuntimeError("Mock engine error") - mock_instance.model_config = MagicMock( - max_model_len=4096, - io_processor_plugin=None, - allowed_local_media_path=None, - allowed_media_domains=None, - ) - # Mimic Qwen3-TTS talker speaker config so CustomVoice validation passes. - mock_instance.model_config.hf_config = MagicMock() - mock_instance.model_config.hf_config.talker_config = MagicMock() - mock_instance.model_config.hf_config.talker_config.speaker_id = {"Vivian": 0} - mock_instance.io_processor = MagicMock() - mock_instance.input_processor = MagicMock() - mock_instance.shutdown = MagicMock() - mock_instance.get_vllm_config = AsyncMock(return_value=None) - mock_instance.get_supported_tasks = AsyncMock(return_value=["generate"]) - mock_instance.get_tokenizer = AsyncMock(return_value=None) + mock_instance = mocker.AsyncMock(spec=RealAsyncOmni) + mock_instance.generate = _build_mock_outputs(server_case.outputs, sampling_case, server_case) - MockAsyncOmni.return_value = mock_instance - yield MockAsyncOmni + mock_instance.stage_list = server_case.stage_list + mock_instance.stage_configs = server_case.stage_configs + mock_instance.output_modalities = _build_output_modalities(server_case.stage_configs) + mock_instance.default_sampling_params_list = [ + SamplingParams() if _stage_type(stage) != "diffusion" else mocker.MagicMock() + for stage in server_case.stage_configs + ] + mock_instance.errored = False + mock_instance.dead_error = RuntimeError("Mock engine error") + mock_instance.model_config = mocker.MagicMock( + max_model_len=4096, + io_processor_plugin=None, + allowed_local_media_path=None, + allowed_media_domains=None, + ) + # Mimic Qwen3-TTS talker speaker config so CustomVoice validation passes. + mock_instance.model_config.hf_config = mocker.MagicMock() + mock_instance.model_config.hf_config.talker_config = mocker.MagicMock() + mock_instance.model_config.hf_config.talker_config.speaker_id = {"Vivian": 0} + mock_instance.io_processor = mocker.MagicMock() + mock_instance.input_processor = mocker.MagicMock() + mock_instance.shutdown = mocker.MagicMock() + mock_instance.get_vllm_config = mocker.AsyncMock(return_value=None) + mock_instance.get_supported_tasks = mocker.AsyncMock(return_value=["generate"]) + mock_instance.get_tokenizer = mocker.AsyncMock(return_value=None) + + mock_async_omni_cls.return_value = mock_instance + yield mock_async_omni_cls @pytest.fixture @@ -583,9 +588,9 @@ async def test_image_generation_node(api_server: str, model: str, image_input: b ServerCase( served_model="Qwen/Qwen2.5-Omni-7B", stage_list=[ - MagicMock(is_comprehension=True, model_stage="llm"), - MagicMock(is_comprehension=False, model_stage="llm"), - MagicMock(is_comprehension=False, model_stage="llm"), + SimpleNamespace(is_comprehension=True, model_stage="llm"), + SimpleNamespace(is_comprehension=False, model_stage="llm"), + SimpleNamespace(is_comprehension=False, model_stage="llm"), ], stage_configs=[ _make_stage_config("llm", is_comprehension=True, model_stage="thinker"), diff --git a/tests/core/sched/test_generation_scheduler_restore.py b/tests/core/sched/test_generation_scheduler_restore.py index 154f40b39..5cc1cab70 100644 --- a/tests/core/sched/test_generation_scheduler_restore.py +++ b/tests/core/sched/test_generation_scheduler_restore.py @@ -6,7 +6,6 @@ internal deques but restore_queues() is not called due to an exception, those requests are permanently orphaned. """ -import unittest from collections import deque import pytest @@ -39,7 +38,7 @@ class FakeAdapter: pass -class TestRestoreQueuesOnError(unittest.TestCase): +class TestRestoreQueuesOnError: """Verify that restore_queues is called even when rewrapping raises.""" def test_requests_not_lost_on_exception(self): @@ -52,8 +51,8 @@ class TestRestoreQueuesOnError(unittest.TestCase): # Step 1: process_pending_chunks moves req-B out adapter.process_pending_chunks(waiting=[], running=running) - self.assertEqual(running, ["req-A"]) - self.assertEqual(len(adapter.waiting_for_chunk_running_requests), 1) + assert running == ["req-A"] + assert len(adapter.waiting_for_chunk_running_requests) == 1 # Step 2: simulate the try/except/finally pattern try: @@ -65,9 +64,9 @@ class TestRestoreQueuesOnError(unittest.TestCase): adapter.restore_queues(waiting=[], running=running) # Step 3: verify request is restored - self.assertTrue(adapter.restore_called) - self.assertIn("req-B", running) - self.assertEqual(len(adapter.waiting_for_chunk_running_requests), 0) + assert adapter.restore_called is True + assert "req-B" in running + assert len(adapter.waiting_for_chunk_running_requests) == 0 def test_requests_lost_without_fix(self): """Demonstrate the bug: without restore in except, request is lost.""" @@ -76,7 +75,7 @@ class TestRestoreQueuesOnError(unittest.TestCase): running = ["req-A", "req-B"] adapter.process_pending_chunks(waiting=[], running=running) - self.assertEqual(running, ["req-A"]) + assert running == ["req-A"] # Simulate the BUGGY code: except without restore try: @@ -85,8 +84,8 @@ class TestRestoreQueuesOnError(unittest.TestCase): pass # Bug: no restore_queues call # Request is lost! - self.assertNotIn("req-B", running) - self.assertEqual(len(adapter.waiting_for_chunk_running_requests), 1) + assert "req-B" not in running + assert len(adapter.waiting_for_chunk_running_requests) == 1 def test_happy_path_restores_via_finally(self): """When no exception, restore_queues is still called via finally.""" @@ -102,9 +101,5 @@ class TestRestoreQueuesOnError(unittest.TestCase): finally: adapter.restore_queues(waiting=[], running=running) - self.assertTrue(adapter.restore_called) - self.assertIn("req-B", running) - - -if __name__ == "__main__": - unittest.main() + assert adapter.restore_called is True + assert "req-B" in running diff --git a/tests/diffusion/distributed/test_distributed_vae_executor.py b/tests/diffusion/distributed/test_distributed_vae_executor.py index dc491dcda..b2ee7c10d 100644 --- a/tests/diffusion/distributed/test_distributed_vae_executor.py +++ b/tests/diffusion/distributed/test_distributed_vae_executor.py @@ -1,4 +1,4 @@ -from unittest.mock import MagicMock, patch +from types import SimpleNamespace import pytest import torch @@ -61,40 +61,31 @@ class E2EOperator: class DummyMixin(DistributedVaeMixin): def __init__(self): self.use_tiling = True - self.distributed_executor = MagicMock() - self.distributed_executor.parallel_size = 2 - self.distributed_executor.group = None + self.distributed_executor = SimpleNamespace(parallel_size=2, group=None) @pytest.fixture(autouse=True) -def mock_dist(): - with ( - patch.object(dist, "get_world_size", return_value=2), - patch.object(dist, "get_rank", return_value=0), - patch.object(dist, "is_initialized", return_value=True), - patch.object(dist, "all_reduce", return_value=None), - patch.object(dist, "gather", return_value=None), - patch.object(dist, "broadcast", return_value=None), - ): - yield +def mock_dist(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(dist, "get_world_size", lambda *args, **kwargs: 2) + monkeypatch.setattr(dist, "get_rank", lambda *args, **kwargs: 0) + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "all_reduce", lambda *args, **kwargs: None) + monkeypatch.setattr(dist, "gather", lambda *args, **kwargs: None) + monkeypatch.setattr(dist, "broadcast", lambda *args, **kwargs: None) @pytest.fixture(autouse=True) -def mock_dit_group(): - with patch( +def mock_dit_group(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( "vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_dit_group", - new=MagicMock(return_value=None), - ): - yield + lambda: None, + ) @pytest.fixture(autouse=True) -def mock_dist_vae_executor(): - with ( - patch.object(DistributedVaeExecutor, "gather_tensors", side_effect=lambda x: [x]), - patch.object(DistributedVaeExecutor, "broadcast_tensor", side_effect=lambda x: x), - ): - yield +def mock_dist_vae_executor(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(DistributedVaeExecutor, "gather_tensors", lambda self, x: [x]) + monkeypatch.setattr(DistributedVaeExecutor, "broadcast_tensor", lambda self, x: x) # ============================ diff --git a/tests/diffusion/models/bagel/test_trajectory_recording.py b/tests/diffusion/models/bagel/test_trajectory_recording.py index 80b3f9d9b..345eac107 100644 --- a/tests/diffusion/models/bagel/test_trajectory_recording.py +++ b/tests/diffusion/models/bagel/test_trajectory_recording.py @@ -4,10 +4,10 @@ import types from dataclasses import dataclass -from unittest.mock import MagicMock, patch import pytest import torch +from pytest_mock import MockerFixture from vllm_omni.diffusion.models.bagel.bagel_transformer import ( Bagel, @@ -23,9 +23,9 @@ NUM_TIMESTEPS = 5 EXPECTED_STEPS = NUM_TIMESTEPS - 1 -def _make_mock_bagel(): +def _make_mock_bagel(mocker: MockerFixture): """Create a mock Bagel with forward returning constant velocity.""" - mock = MagicMock(spec=Bagel) + mock = mocker.MagicMock(spec=Bagel) mock._sp_size = 1 # forward returns a small constant velocity so x_t changes each step @@ -78,18 +78,22 @@ def _make_generate_args(num_tokens=NUM_TOKENS, hidden_dim=HIDDEN_DIM, cfg=False) @pytest.fixture(params=[False, True], ids=["no_cfg", "batched_cfg"]) -def bagel_and_args(request): +def bagel_and_args( + request, + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, +): """Mock Bagel instance and generate_image arguments. Parametrized over CFG mode so every test runs on both the no-CFG and batched-CFG code paths. """ cfg = request.param - with patch( + monkeypatch.setattr( "vllm_omni.diffusion.models.bagel.bagel_transformer.get_classifier_free_guidance_world_size", - return_value=1, - ): - yield _make_mock_bagel(), _make_generate_args(cfg=cfg) + lambda: 1, + ) + yield _make_mock_bagel(mocker), _make_generate_args(cfg=cfg) class TestTrajectoryRecording: @@ -188,12 +192,16 @@ class TestTrajectoryLogProbs: """Tests for log-prob recording when a scheduler is provided.""" @pytest.fixture() - def bagel_scheduler_args(self): - with patch( + def bagel_scheduler_args( + self, + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, + ): + monkeypatch.setattr( "vllm_omni.diffusion.models.bagel.bagel_transformer.get_classifier_free_guidance_world_size", - return_value=1, - ): - yield _make_mock_bagel(), _make_generate_args(), _MockScheduler() + lambda: 1, + ) + yield _make_mock_bagel(mocker), _make_generate_args(), _MockScheduler() def test_log_probs_recorded_with_scheduler(self, bagel_scheduler_args): bagel, args, scheduler = bagel_scheduler_args diff --git a/tests/diffusion/models/flux2/test_flux2_transformer_tp.py b/tests/diffusion/models/flux2/test_flux2_transformer_tp.py index faad08afd..54dda1dd0 100644 --- a/tests/diffusion/models/flux2/test_flux2_transformer_tp.py +++ b/tests/diffusion/models/flux2/test_flux2_transformer_tp.py @@ -1,7 +1,6 @@ -from unittest.mock import MagicMock, patch - import pytest import torch +from pytest_mock import MockerFixture from tests.utils import hardware_test from vllm_omni.diffusion.models.flux2.flux2_transformer import ( @@ -12,14 +11,17 @@ from vllm_omni.diffusion.models.flux2.flux2_transformer import ( # Initialize TP group before tests @pytest.fixture(scope="function", autouse=True) -def setup_tp_group(): +def setup_tp_group(mocker: MockerFixture): """Set up TP group for each test function""" - with patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=2): - with patch("vllm.distributed.parallel_state.get_tp_group") as mock_get_tp_group: - mock_tp_group = MagicMock() - mock_tp_group.world_size = 2 - mock_get_tp_group.return_value = mock_tp_group - yield + mocker.patch( + "vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", + return_value=2, + ) + mock_get_tp_group = mocker.patch("vllm.distributed.parallel_state.get_tp_group") + mock_tp_group = mocker.MagicMock() + mock_tp_group.world_size = 2 + mock_get_tp_group.return_value = mock_tp_group + yield class TestFlux2TransformerWeightLoading: diff --git a/tests/diffusion/offloader/test_sequential_backend.py b/tests/diffusion/offloader/test_sequential_backend.py index d18637a78..2539cc068 100644 --- a/tests/diffusion/offloader/test_sequential_backend.py +++ b/tests/diffusion/offloader/test_sequential_backend.py @@ -3,8 +3,6 @@ """Unit tests for SequentialOffloadBackend.""" -from unittest.mock import patch - import pytest import torch from torch import nn @@ -44,7 +42,7 @@ def _track_pin_memory_calls(): class TestMoveParamsPinMemory: - def test_dtensor_skips_pin_memory(self, accelerator_device): + def test_dtensor_skips_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch): """DTensor should skip pin_memory to avoid RuntimeError.""" module = _create_simple_module().to(accelerator_device) tracker, mock_pin = _track_pin_memory_calls() @@ -56,73 +54,73 @@ class TestMoveParamsPinMemory: return True return original_isinstance(obj, cls) - with patch.object(torch.Tensor, "pin_memory", mock_pin): - with patch("builtins.isinstance", fake_isinstance): - hook = SequentialOffloadHook( - offload_targets=[], - device=accelerator_device, - pin_memory=True, - use_hsdp=False, - ) - hook._move_params( - module, - torch.device("cpu"), - non_blocking=False, - pin_memory=True, - ) - assert not tracker["called"], "pin_memory should not be called for DTensor" + monkeypatch.setattr(torch.Tensor, "pin_memory", mock_pin) + monkeypatch.setattr("builtins.isinstance", fake_isinstance) + hook = SequentialOffloadHook( + offload_targets=[], + device=accelerator_device, + pin_memory=True, + use_hsdp=False, + ) + hook._move_params( + module, + torch.device("cpu"), + non_blocking=False, + pin_memory=True, + ) + assert not tracker["called"], "pin_memory should not be called for DTensor" - def test_regular_tensor_calls_pin_memory(self, accelerator_device): + def test_regular_tensor_calls_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch): """Regular tensor should call pin_memory when moving to CPU.""" module = _create_simple_module().to(accelerator_device) tracker, mock_pin = _track_pin_memory_calls() - with patch.object(torch.Tensor, "pin_memory", mock_pin): - hook = SequentialOffloadHook( - offload_targets=[], - device=accelerator_device, - pin_memory=True, - use_hsdp=False, - ) - hook._move_params( - module, - torch.device("cpu"), - non_blocking=False, - pin_memory=True, - ) - assert tracker["called"], "pin_memory should be called for regular tensors" + monkeypatch.setattr(torch.Tensor, "pin_memory", mock_pin) + hook = SequentialOffloadHook( + offload_targets=[], + device=accelerator_device, + pin_memory=True, + use_hsdp=False, + ) + hook._move_params( + module, + torch.device("cpu"), + non_blocking=False, + pin_memory=True, + ) + assert tracker["called"], "pin_memory should be called for regular tensors" - def test_pin_memory_skipped_when_disabled(self, accelerator_device): + def test_pin_memory_skipped_when_disabled(self, accelerator_device, monkeypatch: pytest.MonkeyPatch): """pin_memory should not be called when pin_memory=False.""" module = _create_simple_module().to(accelerator_device) tracker, mock_pin = _track_pin_memory_calls() - with patch.object(torch.Tensor, "pin_memory", mock_pin): - hook = SequentialOffloadHook( - offload_targets=[], - device=accelerator_device, - pin_memory=False, - use_hsdp=False, - ) - hook._move_params( - module, - torch.device("cpu"), - non_blocking=False, - pin_memory=False, - ) - assert not tracker["called"], "pin_memory should not be called when disabled" + monkeypatch.setattr(torch.Tensor, "pin_memory", mock_pin) + hook = SequentialOffloadHook( + offload_targets=[], + device=accelerator_device, + pin_memory=False, + use_hsdp=False, + ) + hook._move_params( + module, + torch.device("cpu"), + non_blocking=False, + pin_memory=False, + ) + assert not tracker["called"], "pin_memory should not be called when disabled" - def test_pin_memory_skipped_for_non_cpu_target(self, accelerator_device): + def test_pin_memory_skipped_for_non_cpu_target(self, accelerator_device, monkeypatch: pytest.MonkeyPatch): """pin_memory should not be called for non-CPU targets.""" module = _create_simple_module().to("cpu") tracker, mock_pin = _track_pin_memory_calls() - with patch.object(torch.Tensor, "pin_memory", mock_pin): - hook = SequentialOffloadHook( - offload_targets=[], - device=torch.device("cpu"), - pin_memory=True, - use_hsdp=False, - ) - hook._move_params(module, accelerator_device, non_blocking=False, pin_memory=True) - assert not tracker["called"], "pin_memory should not be called for non-CPU target" + monkeypatch.setattr(torch.Tensor, "pin_memory", mock_pin) + hook = SequentialOffloadHook( + offload_targets=[], + device=torch.device("cpu"), + pin_memory=True, + use_hsdp=False, + ) + hook._move_params(module, accelerator_device, non_blocking=False, pin_memory=True) + assert not tracker["called"], "pin_memory should not be called for non-CPU target" diff --git a/tests/diffusion/quantization/test_int8_config.py b/tests/diffusion/quantization/test_int8_config.py index d4d5aa5a7..875277ece 100644 --- a/tests/diffusion/quantization/test_int8_config.py +++ b/tests/diffusion/quantization/test_int8_config.py @@ -2,8 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for Int8 quantization config.""" -from unittest.mock import MagicMock, patch - import pytest import torch from pytest_mock import MockerFixture @@ -102,7 +100,7 @@ def test_quantization_config_string_and_dict_equivalent(): assert config_str.quantization_config.activation_scheme == config_dict.quantization_config.activation_scheme -def test_get_quant_method(mocker: MockerFixture): +def test_get_quant_method(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): """Test for get_quant_method method for GPU""" from vllm_omni.quantization.int8_config import Int8OnlineLinearMethod @@ -111,18 +109,16 @@ def test_get_quant_method(mocker: MockerFixture): def _fake_init(self, quant_config): pass - layer = MagicMock(spec=LinearBase) + layer = mocker.Mock(spec=LinearBase) mocker.patch.object(Int8OnlineLinearMethod, "__init__", _fake_init) prefix = "test_layer" # Mock the platform to be GPU - with ( - patch("vllm_omni.platforms.current_omni_platform.is_cuda", return_value=True), - patch("vllm_omni.platforms.current_omni_platform.is_npu", return_value=False), - ): - method = config.get_quant_method(layer, prefix) - assert isinstance(method, Int8OnlineLinearMethod) + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: False) + method = config.get_quant_method(layer, prefix) + assert isinstance(method, Int8OnlineLinearMethod) # Test skipping quantization for a layer config.ignored_layers = [prefix] @@ -130,22 +126,20 @@ def test_get_quant_method(mocker: MockerFixture): assert isinstance(method, UnquantizedLinearMethod) -def test_get_npu_quant_method(): +def test_get_npu_quant_method(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): """Test for get_quant_method method for NPU""" from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod config = build_quant_config("int8") - layer = MagicMock(spec=LinearBase) + layer = mocker.Mock(spec=LinearBase) prefix = "test_layer" # Mock the platform to be NPU - with ( - patch("vllm_omni.platforms.current_omni_platform.is_cuda", return_value=False), - patch("vllm_omni.platforms.current_omni_platform.is_npu", return_value=True), - ): - method = config.get_quant_method(layer, prefix) - assert isinstance(method, NPUInt8OnlineLinearMethod) + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + method = config.get_quant_method(layer, prefix) + assert isinstance(method, NPUInt8OnlineLinearMethod) # Test skipping quantization for a layer config.ignored_layers = [prefix] @@ -245,7 +239,7 @@ class TestNPUInt8LinearMethod: @pytest.fixture def mock_torch_npu(self, mocker): - torch_npu = MagicMock() + torch_npu = mocker.MagicMock() mocker.patch("vllm_omni.quantization.int8_config.torch_npu", return_value=torch_npu) mocker.patch( diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index 4324ba1e6..a64d9920e 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/tests/diffusion/test_diffusion_scheduler.py @@ -4,10 +4,10 @@ import queue import threading from types import SimpleNamespace -from unittest.mock import Mock, patch import pytest import torch +from pytest_mock import MockerFixture from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError from vllm_omni.diffusion.diffusion_engine import DiffusionEngine @@ -97,19 +97,19 @@ class _StubScheduler(SchedulerInterface): def add_request(self, request: OmniDiffusionRequest) -> str: assert request is self._request - self._state = Mock(sched_req_id=self._sched_req_id, req=request) + self._state = SimpleNamespace(sched_req_id=self._sched_req_id, req=request) return self._sched_req_id def schedule(self): if self._scheduled or self._state is None: - return Mock( + return SimpleNamespace( scheduled_new_reqs=[], scheduled_cached_reqs=CachedRequestData.make_empty(), scheduled_req_ids=[], is_empty=True, ) self._scheduled = True - return Mock( + return SimpleNamespace( scheduled_new_reqs=[NewRequestData.from_state(self._state)], scheduled_cached_reqs=CachedRequestData.make_empty(), scheduled_req_ids=[self._state.sched_req_id], @@ -153,7 +153,7 @@ class _StubScheduler(SchedulerInterface): class TestRequestScheduler: def setup_method(self) -> None: self.scheduler: RequestScheduler = RequestScheduler() - self.scheduler.initialize(Mock()) + self.scheduler.initialize(SimpleNamespace()) def test_single_request_success_lifecycle(self) -> None: req_id = self.scheduler.add_request(_make_request("a")) @@ -276,23 +276,23 @@ class TestRequestScheduler: class TestDiffusionEngine: - def test_add_req_and_wait_for_response_single_path(self) -> None: + def test_add_req_and_wait_for_response_single_path(self, mocker: MockerFixture) -> None: engine = DiffusionEngine.__new__(DiffusionEngine) engine.scheduler = RequestScheduler() - engine.scheduler.initialize(Mock()) + engine.scheduler.initialize(SimpleNamespace()) engine._rpc_lock = threading.RLock() engine.abort_queue = queue.Queue() request = _make_request("engine") runner_output = _make_request_output("engine") - engine.execute_fn = Mock(return_value=runner_output) + engine.execute_fn = mocker.Mock(return_value=runner_output) output = engine.add_req_and_wait_for_response(request) assert output is runner_output.result engine.execute_fn.assert_called_once() - def test_supports_scheduler_interface_injection(self) -> None: + def test_supports_scheduler_interface_injection(self, mocker: MockerFixture) -> None: request = _make_request("engine_iface") runner_output = _make_request_output("engine_iface") scheduler = _StubScheduler(request, runner_output) @@ -301,33 +301,45 @@ class TestDiffusionEngine: engine.scheduler = scheduler engine._rpc_lock = threading.RLock() engine.abort_queue = queue.Queue() - engine.execute_fn = Mock(return_value=runner_output) + engine.execute_fn = mocker.Mock(return_value=runner_output) output = engine.add_req_and_wait_for_response(request) assert output is runner_output.result engine.execute_fn.assert_called_once() - def test_initializes_injected_scheduler(self) -> None: + def test_initializes_injected_scheduler( + self, + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, + ) -> None: request = _make_request("init") scheduler = _StubScheduler(request, DiffusionOutput(output=None)) - od_config = Mock(model_class_name="mock_model") - fake_executor_cls = Mock(return_value=Mock()) + od_config = SimpleNamespace(model_class_name="mock_model") + fake_executor_cls = mocker.Mock(return_value=mocker.Mock()) - with ( - patch("vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func", return_value=None), - patch("vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func", return_value=None), - patch("vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class", return_value=fake_executor_cls), - patch.object(DiffusionEngine, "_dummy_run", return_value=None), - ): - DiffusionEngine(od_config, scheduler=scheduler) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class", + lambda *args, **kwargs: fake_executor_cls, + ) + monkeypatch.setattr(DiffusionEngine, "_dummy_run", lambda self: None) + + DiffusionEngine(od_config, scheduler=scheduler) assert scheduler.initialized_with is od_config fake_executor_cls.assert_called_once_with(od_config) def test_scheduler_alias_keeps_default_request_scheduler(self) -> None: scheduler = Scheduler() - scheduler.initialize(Mock()) + scheduler.initialize(SimpleNamespace()) req_id = scheduler.add_request(_make_request("alias")) sched_output = scheduler.schedule() @@ -336,10 +348,10 @@ class TestDiffusionEngine: assert req_id in finished assert scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED - def test_step_raises_aborted_error(self) -> None: + def test_step_raises_aborted_error(self, mocker: MockerFixture) -> None: engine = DiffusionEngine.__new__(DiffusionEngine) engine.pre_process_func = None - engine.add_req_and_wait_for_response = Mock( + engine.add_req_and_wait_for_response = mocker.Mock( return_value=DiffusionOutput(aborted=True, abort_message="Request req-abort aborted.") ) @@ -349,7 +361,7 @@ class TestDiffusionEngine: def test_abort_queue_marks_request_finished_aborted(self) -> None: engine = DiffusionEngine.__new__(DiffusionEngine) engine.scheduler = RequestScheduler() - engine.scheduler.initialize(Mock()) + engine.scheduler.initialize(SimpleNamespace()) engine.abort_queue = queue.Queue() req_id = engine.scheduler.add_request(_make_request("req-abort")) @@ -361,7 +373,7 @@ class TestDiffusionEngine: def test_finalize_finished_request_returns_aborted_output(self) -> None: engine = DiffusionEngine.__new__(DiffusionEngine) engine.scheduler = RequestScheduler() - engine.scheduler.initialize(Mock()) + engine.scheduler.initialize(SimpleNamespace()) req_id = engine.scheduler.add_request(_make_request("req-finalize")) engine.scheduler.finish_requests(req_id, DiffusionRequestStatus.FINISHED_ABORTED) @@ -371,29 +383,40 @@ class TestDiffusionEngine: assert output.aborted is True assert output.abort_message == "Request req-finalize aborted." - def test_initializes_step_scheduler_when_step_execution_enabled(self) -> None: - od_config = Mock(model_class_name="mock_model") + def test_initializes_step_scheduler_when_step_execution_enabled( + self, + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, + ) -> None: + od_config = SimpleNamespace(model_class_name="mock_model") od_config.step_execution = True - fake_executor = Mock() - fake_executor_cls = Mock(return_value=fake_executor) + fake_executor = mocker.Mock() + fake_executor_cls = mocker.Mock(return_value=fake_executor) - with ( - patch("vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func", return_value=None), - patch("vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func", return_value=None), - patch("vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class", return_value=fake_executor_cls), - patch.object(DiffusionEngine, "_dummy_run", return_value=None), - ): - engine = DiffusionEngine(od_config) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class", + lambda *args, **kwargs: fake_executor_cls, + ) + monkeypatch.setattr(DiffusionEngine, "_dummy_run", lambda self: None) + engine = DiffusionEngine(od_config) assert isinstance(engine.scheduler, StepScheduler) assert engine.execute_fn is fake_executor.execute_step fake_executor_cls.assert_called_once_with(od_config) - def test_dummy_run_raises_on_output_error(self) -> None: + def test_dummy_run_raises_on_output_error(self, mocker: MockerFixture) -> None: engine = DiffusionEngine.__new__(DiffusionEngine) - engine.od_config = Mock(model_class_name="mock_model") + engine.od_config = SimpleNamespace(model_class_name="mock_model") engine.pre_process_func = None - engine.add_req_and_wait_for_response = Mock(return_value=DiffusionOutput(error="boom")) + engine.add_req_and_wait_for_response = mocker.Mock(return_value=DiffusionOutput(error="boom")) with pytest.raises(RuntimeError, match="Dummy run failed: boom"): engine._dummy_run() @@ -402,7 +425,7 @@ class TestDiffusionEngine: class TestStepScheduler: def setup_method(self) -> None: self.scheduler: StepScheduler = StepScheduler() - self.scheduler.initialize(Mock()) + self.scheduler.initialize(SimpleNamespace()) def test_single_request_step_lifecycle(self) -> None: request = _make_step_request("step", num_inference_steps=3) diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index 68aba9ba3..42687d4a1 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/tests/diffusion/test_diffusion_step_pipeline.py @@ -7,10 +7,10 @@ import queue import threading from contextlib import contextmanager from types import SimpleNamespace -from unittest.mock import Mock import pytest import torch +from pytest_mock import MockerFixture import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module from tests.utils import hardware_test @@ -542,11 +542,11 @@ class TestWorker: class TestExecutor: """MultiprocDiffusionExecutor.execute_step""" - def test_execute_step_passes_through_runner_output(self): + def test_execute_step_passes_through_runner_output(self, mocker: MockerFixture): executor = object.__new__(MultiprocDiffusionExecutor) executor._ensure_open = lambda: None expected = RunnerOutput(req_id="req-step", step_index=1, finished=False, result=None) - executor.collective_rpc = Mock(return_value=expected) + executor.collective_rpc = mocker.Mock(return_value=expected) request = _make_engine_request("req-step", num_inference_steps=2) scheduler_output = _make_scheduler_output(request, sched_req_id="req-step") @@ -578,9 +578,9 @@ class TestEngine: ), ], ) - def test_step_engine_returns_error(self, execute_fn, expected_error): + def test_step_engine_returns_error(self, execute_fn, expected_error, mocker: MockerFixture): scheduler = StepScheduler() - scheduler.initialize(Mock()) + scheduler.initialize(mocker.Mock()) engine = _make_engine(scheduler, execute_fn=execute_fn) output = engine.add_req_and_wait_for_response(_make_engine_request("req-error", num_inference_steps=2)) @@ -588,9 +588,9 @@ class TestEngine: assert output.output is None assert expected_error in output.error - def test_step_execution_completes(self): + def test_step_execution_completes(self, mocker: MockerFixture): scheduler = StepScheduler() - scheduler.initialize(Mock()) + scheduler.initialize(mocker.Mock()) engine = _make_engine(scheduler) request = _make_engine_request("req-step", num_inference_steps=2) @@ -614,9 +614,9 @@ class TestEngine: assert output.error is None assert torch.equal(output.output, torch.tensor([2.0])) - def test_step_abort_stops_rescheduling_after_first_step(self): + def test_step_abort_stops_rescheduling_after_first_step(self, mocker: MockerFixture): scheduler = StepScheduler() - scheduler.initialize(Mock()) + scheduler.initialize(mocker.Mock()) engine = _make_engine(scheduler) request = _make_engine_request("req-stop", num_inference_steps=4) @@ -639,9 +639,9 @@ class TestEngine: assert step["n"] == 1 _assert_aborted_output(output, "req-stop") - def test_step_abort_after_reschedule_returns_aborted_output(self): + def test_step_abort_after_reschedule_returns_aborted_output(self, mocker: MockerFixture): scheduler = StepScheduler() - scheduler.initialize(Mock()) + scheduler.initialize(mocker.Mock()) engine = _make_engine(scheduler) request = _make_engine_request("req-mid", num_inference_steps=4) @@ -666,9 +666,9 @@ class TestEngine: assert step["n"] == 2 _assert_aborted_output(output, "req-mid") - def test_finished_step_without_result_returns_error(self): + def test_finished_step_without_result_returns_error(self, mocker: MockerFixture): scheduler = StepScheduler() - scheduler.initialize(Mock()) + scheduler.initialize(mocker.Mock()) engine = _make_engine( scheduler, execute_fn=lambda _: RunnerOutput( diff --git a/tests/diffusion/test_diffusion_worker_cuda_profiler.py b/tests/diffusion/test_diffusion_worker_cuda_profiler.py index ddc2aed2f..4a3b22c21 100644 --- a/tests/diffusion/test_diffusion_worker_cuda_profiler.py +++ b/tests/diffusion/test_diffusion_worker_cuda_profiler.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from unittest.mock import MagicMock - import pytest from pytest_mock import MockerFixture @@ -55,8 +53,8 @@ class TestDiffusionWorkerCudaProfiler: mock_diffusion_worker_dependencies, ): fake_profiler = mocker.Mock() - fake_profiler.start = MagicMock() - fake_profiler.stop = MagicMock() + fake_profiler.start = mocker.Mock() + fake_profiler.stop = mocker.Mock() mocker.patch( "vllm_omni.diffusion.worker.diffusion_worker.CudaProfilerWrapper", return_value=fake_profiler, diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 517f98dda..4bc3e05fe 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -3,7 +3,7 @@ import queue import threading -from unittest.mock import Mock, patch +from types import SimpleNamespace import pytest import torch @@ -24,11 +24,9 @@ def _tagged_output(tag: str) -> DiffusionOutput: return DiffusionOutput(output=torch.tensor([0]), error=tag) -def _mock_request(tag: str) -> Mock: - """Return a mock ``OmniDiffusionRequest`` identifiable by *tag*.""" - req = Mock() - req.request_ids = [tag] - return req +def _mock_request(tag: str): + """Return a lightweight request object identifiable by *tag*.""" + return SimpleNamespace(request_ids=[tag]) def _make_executor(num_gpus: int = 1): @@ -36,20 +34,18 @@ def _make_executor(num_gpus: int = 1): Returns ``(executor, request_queue, result_queue)``. """ - od_cfg = Mock() - od_cfg.num_gpus = num_gpus - - with patch.object(MultiprocDiffusionExecutor, "_init_executor"): - executor = MultiprocDiffusionExecutor(od_cfg) + od_cfg = SimpleNamespace(num_gpus=num_gpus) + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr(MultiprocDiffusionExecutor, "_init_executor", lambda self: None) + executor = MultiprocDiffusionExecutor(od_cfg) + monkeypatch.undo() req_q: queue.Queue = queue.Queue() res_q: queue.Queue = queue.Queue() - mock_broadcast_mq = Mock() - mock_broadcast_mq.enqueue = req_q.put + mock_broadcast_mq = SimpleNamespace(enqueue=req_q.put) - mock_rmq = Mock() - mock_rmq.dequeue = lambda timeout=None: res_q.get(timeout=timeout if timeout is not None else 10) + mock_rmq = SimpleNamespace(dequeue=lambda timeout=None: res_q.get(timeout=timeout if timeout is not None else 10)) executor._broadcast_mq = mock_broadcast_mq executor._result_mq = mock_rmq @@ -63,7 +59,7 @@ def _make_engine(num_gpus: int = 1): executor, req_q, res_q = _make_executor(num_gpus) engine = DiffusionEngine.__new__(DiffusionEngine) sched = RequestScheduler() - sched.initialize(Mock()) + sched.initialize(SimpleNamespace()) engine.scheduler = sched engine.executor = executor engine._rpc_lock = threading.RLock() diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index cb1f31164..a1fc18f84 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -6,7 +6,7 @@ explicitly patch values that differ from vLLM. import argparse import inspect -from unittest.mock import Mock +from types import SimpleNamespace import pytest from pydantic import ValidationError @@ -102,7 +102,7 @@ def test_qwen3_tts_codec_frame_rate_patching(): vllm_config = EngineArgs().create_model_config() # Create a mock talking config with a dummy value for position_id_per_seconds - mock_talker_config = Mock() + mock_talker_config = SimpleNamespace() mock_talker_config.position_id_per_seconds = 12.3 vllm_config.hf_config.talker_config = mock_talker_config @@ -146,13 +146,12 @@ def test_stage_specific_text_config_override(): # Switch the created hf text config with a mock whose # values we want to pull through the text config helper stage_text_config = vllm_config.hf_text_config - vllm_config.hf_text_config = Mock() + vllm_config.hf_text_config = SimpleNamespace() stage_text_config.sliding_window = 4096 stage_text_config.attention_chunk_size = 2048 # Move the stage config's text config getter & thinker config - mock_stage_config = Mock() - mock_stage_config.get_text_config.return_value = stage_text_config + mock_stage_config = SimpleNamespace(get_text_config=lambda: stage_text_config) vllm_config.hf_config.thinker_config = mock_stage_config # Ensure that create from a vLLM config correctly pulls the diff --git a/tests/engine/test_async_omni_engine_input.py b/tests/engine/test_async_omni_engine_input.py index ed6a7277b..3700e426d 100644 --- a/tests/engine/test_async_omni_engine_input.py +++ b/tests/engine/test_async_omni_engine_input.py @@ -1,6 +1,5 @@ -from unittest.mock import Mock - import pytest +from pytest_mock import MockerFixture from vllm.sampling_params import SamplingParams from vllm.v1.engine import EngineCoreRequest @@ -24,18 +23,18 @@ def _make_engine_core_request() -> EngineCoreRequest: ) -def test_build_add_request_message_preserves_additional_information(): +def test_build_add_request_message_preserves_additional_information(mocker: MockerFixture): engine = object.__new__(AsyncOmniEngine) params = SamplingParams(max_tokens=8) engine.default_sampling_params_list = [params] engine.stage_metadata = [{"stage_type": "llm"}] engine.supported_tasks = ("speech",) - input_processor = Mock() + input_processor = mocker.Mock() input_processor.process_inputs.return_value = _make_engine_core_request() engine.input_processor = input_processor - output_processor = Mock() + output_processor = mocker.Mock() engine.output_processors = [output_processor] prompt = { @@ -63,18 +62,18 @@ def test_build_add_request_message_preserves_additional_information(): output_processor.add_request.assert_called_once() -def test_build_add_request_message_with_resumable_streaming(): +def test_build_add_request_message_with_resumable_streaming(mocker: MockerFixture): engine = object.__new__(AsyncOmniEngine) params = SamplingParams(max_tokens=8) engine.default_sampling_params_list = [params] engine.stage_metadata = [{"stage_type": "llm"}] engine.supported_tasks = ("generate",) - input_processor = Mock() + input_processor = mocker.Mock() input_processor.process_inputs.return_value = _make_engine_core_request() engine.input_processor = input_processor - output_processor = Mock() + output_processor = mocker.Mock() engine.output_processors = [output_processor] msg = engine._build_add_request_message( diff --git a/tests/engine/test_async_omni_engine_outputs.py b/tests/engine/test_async_omni_engine_outputs.py index ccf9e8cb6..ef3cfab3b 100644 --- a/tests/engine/test_async_omni_engine_outputs.py +++ b/tests/engine/test_async_omni_engine_outputs.py @@ -5,36 +5,36 @@ subsequent attempts to collect output raise RuntimeError. """ import queue -from unittest.mock import MagicMock import pytest +from pytest_mock import MockerFixture from vllm_omni.engine.async_omni_engine import AsyncOmniEngine pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -def _make_engine(output_queue, *, thread_alive: bool = True) -> AsyncOmniEngine: +def _make_engine(output_queue, mocker: MockerFixture, *, thread_alive: bool = True) -> AsyncOmniEngine: """Create an AsyncOmniEngine bypassing __init__.""" engine = object.__new__(AsyncOmniEngine) engine.output_queue = output_queue - engine.orchestrator_thread = MagicMock( - is_alive=MagicMock(return_value=thread_alive), + engine.orchestrator_thread = mocker.MagicMock( + is_alive=mocker.MagicMock(return_value=thread_alive), ) return engine -def test_try_get_output_raises_after_orchestrator_dies(): +def test_try_get_output_raises_after_orchestrator_dies(mocker: MockerFixture): """Draining remaining results then hitting an empty queue with a dead orchestrator must raise RuntimeError so callers know the pipeline is gone.""" - mock_queue = MagicMock() + mock_queue = mocker.MagicMock() # First call succeeds; second call finds the queue empty. mock_queue.sync_q.get.side_effect = [ {"type": "output", "request_id": "r1"}, queue.Empty, ] - engine = _make_engine(mock_queue, thread_alive=True) + engine = _make_engine(mock_queue, mocker, thread_alive=True) # Collect the one buffered result. assert engine.try_get_output()["request_id"] == "r1" @@ -47,15 +47,15 @@ def test_try_get_output_raises_after_orchestrator_dies(): @pytest.mark.asyncio -async def test_try_get_output_async_raises_after_orchestrator_dies(): +async def test_try_get_output_async_raises_after_orchestrator_dies(mocker: MockerFixture): """Same scenario as above but for the async variant.""" - mock_queue = MagicMock() + mock_queue = mocker.MagicMock() mock_queue.sync_q.get_nowait.side_effect = [ {"type": "output", "request_id": "r1"}, queue.Empty, ] - engine = _make_engine(mock_queue, thread_alive=True) + engine = _make_engine(mock_queue, mocker, thread_alive=True) assert (await engine.try_get_output_async())["request_id"] == "r1" diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 2c5bf6cc7..608e92ac4 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -17,10 +17,11 @@ from __future__ import annotations import threading from contextlib import contextmanager +from types import SimpleNamespace from typing import Any -from unittest.mock import MagicMock, Mock, patch import pytest +from pytest_mock import MockerFixture from vllm.v1.engine.utils import EngineZmqAddresses from vllm_omni.engine.async_omni_engine import AsyncOmniEngine @@ -41,31 +42,33 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] # --------------------------------------------------------------------------- -def _make_stage_cfg(stage_id: int, stage_type: str = "llm") -> Mock: +def _make_stage_cfg(stage_id: int, stage_type: str = "llm"): """Return a lightweight stage config mock.""" - cfg = Mock() - cfg.stage_id = stage_id - cfg.stage_type = stage_type - cfg.engine_args = MagicMock() - cfg.engine_args.async_chunk = False - cfg.engine_args.model_stage = None - cfg.engine_args.engine_output_type = None - return cfg + return SimpleNamespace( + stage_id=stage_id, + stage_type=stage_type, + engine_args=SimpleNamespace( + async_chunk=False, + model_stage=None, + engine_output_type=None, + ), + ) def _make_started_llm_stage(stage_id: int) -> StartedLlmStage: """Return a minimal StartedLlmStage for mocking.""" - addresses = Mock() - addresses.inputs = ["tcp://127.0.0.1:5000"] - addresses.outputs = ["tcp://127.0.0.1:5001"] - addresses.frontend_stats_publish_address = None + addresses = SimpleNamespace( + inputs=["tcp://127.0.0.1:5000"], + outputs=["tcp://127.0.0.1:5001"], + frontend_stats_publish_address=None, + ) return StartedLlmStage( stage_id=stage_id, - metadata=Mock(stage_id=stage_id), - vllm_config=Mock(), - executor_class=Mock(), - engine_manager=Mock(), - coordinator=Mock(), + metadata=SimpleNamespace(stage_id=stage_id), + vllm_config=SimpleNamespace(), + executor_class=SimpleNamespace(), + engine_manager=SimpleNamespace(), + coordinator=SimpleNamespace(), addresses=addresses, ) @@ -348,74 +351,80 @@ class TestSingleStageModeDetection: the orchestrator thread, so no actual engines are started. """ - def _make_engine_no_thread(self, **kwargs: Any) -> AsyncOmniEngine: + def _make_engine_no_thread(self, mocker: MockerFixture, **kwargs: Any) -> AsyncOmniEngine: """Create an AsyncOmniEngine without starting the orchestrator thread.""" stage_cfg = _make_stage_cfg(0) mock_stage_configs = [stage_cfg] - with ( - patch.object( - AsyncOmniEngine, - "_resolve_stage_configs", - return_value=("/fake/path", mock_stage_configs), - ), - patch.object( - AsyncOmniEngine, - "_bootstrap_orchestrator", - ), - patch("threading.Thread") as mock_thread_cls, - patch("concurrent.futures.Future") as mock_future_cls, - ): - mock_future = Mock() - mock_future.result.return_value = Mock() # simulates a loop - mock_future_cls.return_value = mock_future + mocker.patch.object( + AsyncOmniEngine, + "_resolve_stage_configs", + return_value=("/fake/path", mock_stage_configs), + ) + mocker.patch.object( + AsyncOmniEngine, + "_bootstrap_orchestrator", + ) + mock_thread_cls = mocker.patch("threading.Thread") + mock_future_cls = mocker.patch("concurrent.futures.Future") - mock_thread = Mock() - mock_thread.is_alive.return_value = False - mock_thread_cls.return_value = mock_thread + mock_future = mocker.Mock() + mock_future.result.return_value = mocker.Mock() # simulates a loop + mock_future_cls.return_value = mock_future - engine = AsyncOmniEngine(model="fake-model", **kwargs) + mock_thread = mocker.Mock() + mock_thread.is_alive.return_value = False + mock_thread_cls.return_value = mock_thread + + engine = AsyncOmniEngine(model="fake-model", **kwargs) return engine - def test_explicit_single_stage_mode_true(self): + def test_explicit_single_stage_mode_true(self, mocker: MockerFixture): engine = self._make_engine_no_thread( + mocker, single_stage_mode=True, omni_master_address="127.0.0.1", omni_master_port=20000, ) assert engine.single_stage_mode is True - def test_stage_id_kwarg_promotes_to_single_stage_mode(self): + def test_stage_id_kwarg_promotes_to_single_stage_mode(self, mocker: MockerFixture): engine = self._make_engine_no_thread( + mocker, stage_id=0, omni_master_address="127.0.0.1", omni_master_port=20001, ) assert engine.single_stage_mode is True - def test_stage_id_kwarg_sets_filter(self): + def test_stage_id_kwarg_sets_filter(self, mocker: MockerFixture): engine = self._make_engine_no_thread( + mocker, stage_id=1, omni_master_address="127.0.0.1", omni_master_port=20002, ) assert engine._single_stage_id_filter == 1 - def test_no_stage_id_no_single_stage_mode(self): - engine = self._make_engine_no_thread() + def test_no_stage_id_no_single_stage_mode(self, mocker: MockerFixture): + engine = self._make_engine_no_thread( + mocker, + ) assert engine.single_stage_mode is False assert engine._single_stage_id_filter is None - def test_single_stage_mode_without_stage_id_has_no_filter(self): + def test_single_stage_mode_without_stage_id_has_no_filter(self, mocker: MockerFixture): engine = self._make_engine_no_thread( + mocker, single_stage_mode=True, omni_master_address="127.0.0.1", omni_master_port=20003, ) assert engine._single_stage_id_filter is None - def test_master_address_and_port_stored(self): + def test_master_address_and_port_stored(self, mocker: MockerFixture): engine = self._make_engine_no_thread( + mocker, stage_id=0, omni_master_address="10.0.0.1", omni_master_port=12345, @@ -423,8 +432,10 @@ class TestSingleStageModeDetection: assert engine._omni_master_address == "10.0.0.1" assert engine._omni_master_port == 12345 - def test_omni_master_server_starts_as_none(self): - engine = self._make_engine_no_thread() + def test_omni_master_server_starts_as_none(self, mocker: MockerFixture): + engine = self._make_engine_no_thread( + mocker, + ) assert engine._omni_master_server is None @@ -448,7 +459,7 @@ class TestInitializeStagesRouting: def _build_engine_skeleton( self, - stage_cfgs: list[Mock], + stage_cfgs: list[Any], single_stage_mode: bool, stage_id_filter: int | None, omni_master_address: str = "127.0.0.1", @@ -478,8 +489,8 @@ class TestInitializeStagesRouting: engine.prompt_expand_func = None return engine - def _fake_metadata(self, stage_id: int, stage_type: str = "llm") -> Mock: - meta = Mock() + def _fake_metadata(self, mocker: MockerFixture, stage_id: int, stage_type: str = "llm") -> Any: + meta = mocker.Mock() meta.stage_id = stage_id meta.stage_type = stage_type meta.runtime_cfg = {} @@ -492,13 +503,14 @@ class TestInitializeStagesRouting: def _run_initialize_stages_mocked( self, + mocker: MockerFixture, engine: AsyncOmniEngine, - stage_cfgs: list[Mock], + stage_cfgs: list[Any], *, launch_side_effect: Any = None, remote_side_effect: Any = None, attach_result: Any = None, - ) -> tuple[Mock, Mock]: + ) -> tuple[Any, Any]: """Execute _initialize_stages with all heavy helpers mocked. Returns (mock_launch_llm_stage, mock_create_remote_llm_stage). @@ -509,167 +521,217 @@ class TestInitializeStagesRouting: if getattr(cfg, "stage_type", "llm") != "diffusion" } - default_attach = (Mock(), Mock(), Mock(), Mock()) + default_attach = (mocker.Mock(), mocker.Mock(), mocker.Mock(), mocker.Mock()) - mock_launch = Mock( + mock_launch = mocker.Mock( side_effect=launch_side_effect or (lambda cfg, meta, spec, timeout, llm_stage_launch_lock, kv: started_by_stage[meta.stage_id]) ) - mock_remote = Mock( + mock_remote = mocker.Mock( side_effect=remote_side_effect or (lambda cfg, meta, spec, timeout, srv: started_by_stage[meta.stage_id]) ) - mock_attach = Mock(return_value=attach_result or default_attach) + mock_attach = mocker.Mock(return_value=attach_result or default_attach) - mock_oms = Mock(spec=OmniMasterServer) - mock_oms.get_zmq_addresses.side_effect = lambda sid: Mock() + mock_oms = mocker.Mock(spec=OmniMasterServer) + mock_oms.get_zmq_addresses.side_effect = lambda sid: mocker.Mock() finalized = ( - [Mock() for _ in stage_cfgs], - [Mock() for _ in stage_cfgs], + [mocker.Mock() for _ in stage_cfgs], + [mocker.Mock() for _ in stage_cfgs], [{"final_output": True, "final_output_type": None, "stage_type": "llm"} for _ in stage_cfgs], ) - with ( - patch.object(engine, "_launch_llm_stage", mock_launch), - patch.object(engine, "_create_remote_llm_stage", mock_remote), - patch.object(engine, "_attach_llm_stage", mock_attach), - patch("vllm_omni.engine.async_omni_engine.OmniMasterServer", return_value=mock_oms), - patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), - patch( - "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", - return_value=None, + mocker.patch.object(engine, "_launch_llm_stage", mock_launch) + mocker.patch.object(engine, "_create_remote_llm_stage", mock_remote) + mocker.patch.object(engine, "_attach_llm_stage", mock_attach) + mocker.patch( + "vllm_omni.engine.async_omni_engine.OmniMasterServer", + return_value=mock_oms, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.prepare_engine_environment", + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", + return_value=None, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", + return_value={}, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.extract_stage_metadata", + side_effect=lambda cfg: self._fake_metadata( + mocker, + cfg.stage_id, + getattr(cfg, "stage_type", "llm"), ), - patch( - "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", - return_value={}, - ), - patch( - "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", - return_value=(None, None, None), - ), - patch( - "vllm_omni.engine.async_omni_engine.extract_stage_metadata", - side_effect=lambda cfg: self._fake_metadata(cfg.stage_id, getattr(cfg, "stage_type", "llm")), - ), - patch( - "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", - return_value=finalized, - ), - ): - engine._initialize_stages(stage_init_timeout=60) + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", + return_value=finalized, + ) + + engine._initialize_stages(stage_init_timeout=60) return mock_launch, mock_remote # -- single-stage mode: stage matches filter → local launch --------------- - def test_matching_stage_uses_launch_llm_stage(self): + def test_matching_stage_uses_launch_llm_stage(self, mocker: MockerFixture): """stage_id == _single_stage_id_filter → _launch_llm_stage is called.""" stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1)] engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=0) - mock_launch, mock_remote = self._run_initialize_stages_mocked(engine, stage_cfgs) + mock_launch, mock_remote = self._run_initialize_stages_mocked(mocker, engine, stage_cfgs) launched_ids = [c.args[1].stage_id for c in mock_launch.call_args_list] assert 0 in launched_ids, "_launch_llm_stage should be called for stage 0" - def test_non_matching_stage_uses_create_remote_llm_stage(self): + def test_non_matching_stage_uses_create_remote_llm_stage(self, mocker: MockerFixture): """stage_id != _single_stage_id_filter → _create_remote_llm_stage is called.""" stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1)] engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=0) - mock_launch, mock_remote = self._run_initialize_stages_mocked(engine, stage_cfgs) + mock_launch, mock_remote = self._run_initialize_stages_mocked(mocker, engine, stage_cfgs) remote_ids = [c.args[1].stage_id for c in mock_remote.call_args_list] assert 1 in remote_ids, "_create_remote_llm_stage should be called for stage 1" - def test_filter_1_routes_correctly(self): + def test_filter_1_routes_correctly(self, mocker: MockerFixture): """With filter=1, stage 0 is remote and stage 1 is local.""" stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1)] engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=1) - mock_launch, mock_remote = self._run_initialize_stages_mocked(engine, stage_cfgs) + mock_launch, mock_remote = self._run_initialize_stages_mocked(mocker, engine, stage_cfgs) launched_ids = [c.args[1].stage_id for c in mock_launch.call_args_list] remote_ids = [c.args[1].stage_id for c in mock_remote.call_args_list] assert 1 in launched_ids, "stage 1 should be launched locally with filter=1" assert 0 in remote_ids, "stage 0 should use remote path with filter=1" - def test_no_filter_all_stages_use_launch_path(self): + def test_no_filter_all_stages_use_launch_path(self, mocker: MockerFixture): """single_stage_mode=True but no filter → all stages use _launch_llm_stage.""" stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1)] engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=None) - mock_launch, mock_remote = self._run_initialize_stages_mocked(engine, stage_cfgs) + mock_launch, mock_remote = self._run_initialize_stages_mocked(mocker, engine, stage_cfgs) assert mock_remote.call_count == 0, "No remote launches without a filter" launched_ids = [c.args[1].stage_id for c in mock_launch.call_args_list] assert set(launched_ids) == {0, 1} - def test_non_single_stage_mode_never_calls_create_remote(self): + def test_non_single_stage_mode_never_calls_create_remote(self, mocker: MockerFixture): """Outside single_stage_mode, _create_remote_llm_stage must not be called.""" stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1)] engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=False, stage_id_filter=None) - mock_launch, mock_remote = self._run_initialize_stages_mocked(engine, stage_cfgs) + mock_launch, mock_remote = self._run_initialize_stages_mocked(mocker, engine, stage_cfgs) assert mock_remote.call_count == 0 - def test_omni_master_server_started_in_single_stage_mode(self): + def test_omni_master_server_started_in_single_stage_mode(self, mocker: MockerFixture): """OmniMasterServer.start() must be called when single_stage_mode=True.""" stage_cfgs = [_make_stage_cfg(0)] engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=0) - mock_oms = Mock(spec=OmniMasterServer) - mock_oms.get_zmq_addresses.return_value = Mock() - finalized = ([Mock()], [Mock()], [{"final_output": True, "final_output_type": None, "stage_type": "llm"}]) + mock_oms = mocker.Mock(spec=OmniMasterServer) + mock_oms.get_zmq_addresses.return_value = mocker.Mock() + finalized = ( + [mocker.Mock()], + [mocker.Mock()], + [{"final_output": True, "final_output_type": None, "stage_type": "llm"}], + ) - with ( - patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(0)), - patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(0)), - patch.object(engine, "_attach_llm_stage", return_value=(Mock(), Mock(), Mock(), Mock())), - patch("vllm_omni.engine.async_omni_engine.OmniMasterServer", return_value=mock_oms), - patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), - patch("vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", return_value=None), - patch("vllm_omni.engine.async_omni_engine.get_stage_connector_spec", return_value={}), - patch( - "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", return_value=(None, None, None) - ), - patch( - "vllm_omni.engine.async_omni_engine.extract_stage_metadata", - side_effect=lambda cfg: self._fake_metadata(cfg.stage_id), - ), - patch("vllm_omni.engine.async_omni_engine.finalize_initialized_stages", return_value=finalized), - ): - engine._initialize_stages(stage_init_timeout=60) + mocker.patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(0)) + mocker.patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(0)) + mocker.patch.object( + engine, + "_attach_llm_stage", + return_value=(mocker.Mock(), mocker.Mock(), mocker.Mock(), mocker.Mock()), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.OmniMasterServer", + return_value=mock_oms, + ) + mocker.patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment") + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", + return_value=None, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", + return_value={}, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.extract_stage_metadata", + side_effect=lambda cfg: self._fake_metadata(mocker, cfg.stage_id), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", + return_value=finalized, + ) + + engine._initialize_stages(stage_init_timeout=60) mock_oms.start.assert_called_once() - def test_omni_master_server_uses_configured_stage_ids(self): + def test_omni_master_server_uses_configured_stage_ids(self, mocker: MockerFixture): """Configured stage IDs, not list indexes, should drive pre-allocation.""" stage_cfgs = [_make_stage_cfg(7), _make_stage_cfg(11)] engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=7) - mock_oms = Mock(spec=OmniMasterServer) - mock_oms.get_zmq_addresses.return_value = Mock() + mock_oms = mocker.Mock(spec=OmniMasterServer) + mock_oms.get_zmq_addresses.return_value = mocker.Mock() finalized = ( - [Mock(), Mock()], - [Mock(), Mock()], + [mocker.Mock(), mocker.Mock()], + [mocker.Mock(), mocker.Mock()], [{"final_output": False, "final_output_type": None, "stage_type": "llm"} for _ in stage_cfgs], ) - with ( - patch.object( - engine, "_launch_llm_stage", side_effect=[_make_started_llm_stage(7), _make_started_llm_stage(11)] - ), - patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(11)), - patch.object(engine, "_attach_llm_stage", return_value=(Mock(), Mock(), Mock(), Mock())), - patch("vllm_omni.engine.async_omni_engine.OmniMasterServer", return_value=mock_oms) as mock_oms_cls, - patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), - patch("vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", return_value=None), - patch("vllm_omni.engine.async_omni_engine.get_stage_connector_spec", return_value={}), - patch( - "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", return_value=(None, None, None) - ), - patch( - "vllm_omni.engine.async_omni_engine.extract_stage_metadata", - side_effect=lambda cfg: self._fake_metadata(cfg.stage_id), - ), - patch("vllm_omni.engine.async_omni_engine.finalize_initialized_stages", return_value=finalized), - ): - engine._initialize_stages(stage_init_timeout=60) + mocker.patch.object( + engine, + "_launch_llm_stage", + side_effect=[_make_started_llm_stage(7), _make_started_llm_stage(11)], + ) + mocker.patch.object( + engine, + "_create_remote_llm_stage", + return_value=_make_started_llm_stage(11), + ) + mocker.patch.object( + engine, + "_attach_llm_stage", + return_value=(mocker.Mock(), mocker.Mock(), mocker.Mock(), mocker.Mock()), + ) + mock_oms_cls = mocker.patch( + "vllm_omni.engine.async_omni_engine.OmniMasterServer", + return_value=mock_oms, + ) + mocker.patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment") + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", + return_value=None, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", + return_value={}, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.extract_stage_metadata", + side_effect=lambda cfg: self._fake_metadata(mocker, cfg.stage_id), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", + return_value=finalized, + ) + + engine._initialize_stages(stage_init_timeout=60) mock_oms_cls.assert_called_once_with( master_address=engine._omni_master_address, @@ -677,73 +739,121 @@ class TestInitializeStagesRouting: stage_ids=[7, 11], ) - def test_single_stage_filter_uses_configured_stage_ids(self): + def test_single_stage_filter_uses_configured_stage_ids(self, mocker: MockerFixture): """Local/remote dispatch should compare against configured stage IDs.""" stage_cfgs = [_make_stage_cfg(7), _make_stage_cfg(11)] engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=7) - mock_oms = Mock(spec=OmniMasterServer) + mock_oms = mocker.Mock(spec=OmniMasterServer) finalized = ( - [Mock(), Mock()], - [Mock(), Mock()], + [mocker.Mock(), mocker.Mock()], + [mocker.Mock(), mocker.Mock()], [{"final_output": False, "final_output_type": None, "stage_type": "llm"} for _ in stage_cfgs], ) - with ( - patch.object(engine, "_launch_llm_stage", side_effect=[_make_started_llm_stage(7)]) as mock_launch, - patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(11)) as mock_remote, - patch.object(engine, "_attach_llm_stage", return_value=(Mock(), Mock(), Mock(), Mock())), - patch("vllm_omni.engine.async_omni_engine.OmniMasterServer", return_value=mock_oms), - patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), - patch("vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", return_value=None), - patch("vllm_omni.engine.async_omni_engine.get_stage_connector_spec", return_value={}), - patch( - "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", return_value=(None, None, None) - ), - patch( - "vllm_omni.engine.async_omni_engine.extract_stage_metadata", - side_effect=lambda cfg: self._fake_metadata(cfg.stage_id), - ), - patch("vllm_omni.engine.async_omni_engine.finalize_initialized_stages", return_value=finalized), - ): - engine._initialize_stages(stage_init_timeout=60) + mock_launch = mocker.patch.object( + engine, + "_launch_llm_stage", + side_effect=[_make_started_llm_stage(7)], + ) + mock_remote = mocker.patch.object( + engine, + "_create_remote_llm_stage", + return_value=_make_started_llm_stage(11), + ) + mocker.patch.object( + engine, + "_attach_llm_stage", + return_value=(mocker.Mock(), mocker.Mock(), mocker.Mock(), mocker.Mock()), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.OmniMasterServer", + return_value=mock_oms, + ) + mocker.patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment") + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", + return_value=None, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", + return_value={}, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.extract_stage_metadata", + side_effect=lambda cfg: self._fake_metadata(mocker, cfg.stage_id), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", + return_value=finalized, + ) + + engine._initialize_stages(stage_init_timeout=60) assert [call.args[1].stage_id for call in mock_launch.call_args_list] == [7] assert [call.args[1].stage_id for call in mock_remote.call_args_list] == [11] - def test_omni_master_server_preallocates_diffusion_stage_ids(self): + def test_omni_master_server_preallocates_diffusion_stage_ids(self, mocker: MockerFixture): """Diffusion stages should also receive OmniMasterServer allocations.""" stage_cfgs = [_make_stage_cfg(7), _make_stage_cfg(11, stage_type="diffusion")] engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=7) - mock_oms = Mock(spec=OmniMasterServer) + mock_oms = mocker.Mock(spec=OmniMasterServer) finalized = ( - [Mock(), Mock()], - [Mock(), Mock()], + [mocker.Mock(), mocker.Mock()], + [mocker.Mock(), mocker.Mock()], [ {"final_output": False, "final_output_type": None, "stage_type": "llm"}, {"final_output": False, "final_output_type": None, "stage_type": "diffusion"}, ], ) - with ( - patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(7)), - patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(7)), - patch.object(engine, "_launch_diffusion_stage", return_value=Mock()), - patch.object(engine, "_create_remote_diffusion_stage", return_value=Mock()), - patch.object(engine, "_attach_llm_stage", return_value=(Mock(), Mock(), Mock(), Mock())), - patch("vllm_omni.engine.async_omni_engine.OmniMasterServer", return_value=mock_oms) as mock_oms_cls, - patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), - patch("vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", return_value=None), - patch("vllm_omni.engine.async_omni_engine.get_stage_connector_spec", return_value={}), - patch( - "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", return_value=(None, None, None) + mocker.patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(7)) + mocker.patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(7)) + mocker.patch.object(engine, "_launch_diffusion_stage", return_value=mocker.Mock()) + mocker.patch.object( + engine, + "_create_remote_diffusion_stage", + return_value=mocker.Mock(), + ) + mocker.patch.object( + engine, + "_attach_llm_stage", + return_value=(mocker.Mock(), mocker.Mock(), mocker.Mock(), mocker.Mock()), + ) + mock_oms_cls = mocker.patch( + "vllm_omni.engine.async_omni_engine.OmniMasterServer", + return_value=mock_oms, + ) + mocker.patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment") + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", + return_value=None, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", + return_value={}, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.extract_stage_metadata", + side_effect=lambda cfg: self._fake_metadata( + mocker, + cfg.stage_id, + getattr(cfg, "stage_type", "llm"), ), - patch( - "vllm_omni.engine.async_omni_engine.extract_stage_metadata", - side_effect=lambda cfg: self._fake_metadata(cfg.stage_id, getattr(cfg, "stage_type", "llm")), - ), - patch("vllm_omni.engine.async_omni_engine.finalize_initialized_stages", return_value=finalized), - ): - engine._initialize_stages(stage_init_timeout=60) + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", + return_value=finalized, + ) + + engine._initialize_stages(stage_init_timeout=60) mock_oms_cls.assert_called_once_with( master_address=engine._omni_master_address, @@ -751,135 +861,200 @@ class TestInitializeStagesRouting: stage_ids=[7, 11], ) - def test_duplicate_llm_stage_ids_raise(self): + def test_duplicate_llm_stage_ids_raise(self, mocker: MockerFixture): """Duplicate configured LLM stage IDs should fail fast.""" stage_cfgs = [_make_stage_cfg(3), _make_stage_cfg(3)] engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=3) - with ( - patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), - patch("vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", return_value=None), - pytest.raises(ValueError, match="Duplicate stage_id"), - ): + mocker.patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment") + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", + return_value=None, + ) + with pytest.raises(ValueError, match="Duplicate stage_id"): engine._initialize_stages(stage_init_timeout=60) - def test_omni_master_server_not_started_in_normal_mode(self): + def test_omni_master_server_not_started_in_normal_mode(self, mocker: MockerFixture): """OmniMasterServer must NOT be instantiated outside single_stage_mode.""" stage_cfgs = [_make_stage_cfg(0)] engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=False, stage_id_filter=None) - finalized = ([Mock()], [Mock()], [{"final_output": True, "final_output_type": None, "stage_type": "llm"}]) + finalized = ( + [mocker.Mock()], + [mocker.Mock()], + [{"final_output": True, "final_output_type": None, "stage_type": "llm"}], + ) - with ( - patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(0)), - patch.object(engine, "_attach_llm_stage", return_value=(Mock(), Mock(), Mock(), Mock())), - patch("vllm_omni.engine.async_omni_engine.OmniMasterServer") as mock_oms_cls, - patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), - patch("vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", return_value=None), - patch("vllm_omni.engine.async_omni_engine.get_stage_connector_spec", return_value={}), - patch( - "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", return_value=(None, None, None) - ), - patch( - "vllm_omni.engine.async_omni_engine.extract_stage_metadata", - side_effect=lambda cfg: self._fake_metadata(cfg.stage_id), - ), - patch("vllm_omni.engine.async_omni_engine.finalize_initialized_stages", return_value=finalized), - ): - engine._initialize_stages(stage_init_timeout=60) + mocker.patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(0)) + mocker.patch.object( + engine, + "_attach_llm_stage", + return_value=(mocker.Mock(), mocker.Mock(), mocker.Mock(), mocker.Mock()), + ) + mock_oms_cls = mocker.patch("vllm_omni.engine.async_omni_engine.OmniMasterServer") + mocker.patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment") + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", + return_value=None, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", + return_value={}, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.extract_stage_metadata", + side_effect=lambda cfg: self._fake_metadata(mocker, cfg.stage_id), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", + return_value=finalized, + ) + + engine._initialize_stages(stage_init_timeout=60) mock_oms_cls.assert_not_called() - def test_single_stage_mode_missing_master_address_raises(self): + def test_single_stage_mode_missing_master_address_raises(self, mocker: MockerFixture): """single_stage_mode without master address/port raises ValueError.""" stage_cfgs = [_make_stage_cfg(0)] engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=0) engine._omni_master_address = None # missing engine._omni_master_port = None - with ( - patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), - patch("vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", return_value=None), - pytest.raises(ValueError, match="omni_master_address"), - ): + mocker.patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment") + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", + return_value=None, + ) + with pytest.raises(ValueError, match="omni_master_address"): engine._initialize_stages(stage_init_timeout=60) - def test_matching_diffusion_stage_uses_local_registered_launch(self): + def test_matching_diffusion_stage_uses_local_registered_launch(self, mocker: MockerFixture): """A local diffusion stage should use the registered single-stage launch path.""" stage_cfgs = [_make_stage_cfg(0, stage_type="diffusion"), _make_stage_cfg(1)] engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=0) - mock_oms = Mock(spec=OmniMasterServer) - diffusion_client = Mock(stage_type="diffusion") + mock_oms = mocker.Mock(spec=OmniMasterServer) + diffusion_client = mocker.Mock(stage_type="diffusion") finalized = ( - [diffusion_client, Mock()], - [Mock(), Mock()], + [diffusion_client, mocker.Mock()], + [mocker.Mock(), mocker.Mock()], [ {"final_output": False, "final_output_type": None, "stage_type": "diffusion"}, {"final_output": False, "final_output_type": None, "stage_type": "llm"}, ], ) - with ( - patch.object(engine, "_launch_diffusion_stage", return_value=diffusion_client) as mock_local_diff, - patch.object(engine, "_create_remote_diffusion_stage") as mock_remote_diff, - patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(1)), - patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(1)), - patch.object(engine, "_attach_llm_stage", return_value=(Mock(), Mock(), Mock(), Mock())), - patch("vllm_omni.engine.async_omni_engine.OmniMasterServer", return_value=mock_oms), - patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), - patch("vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", return_value=None), - patch("vllm_omni.engine.async_omni_engine.get_stage_connector_spec", return_value={}), - patch( - "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", return_value=(None, None, None) + mock_local_diff = mocker.patch.object( + engine, + "_launch_diffusion_stage", + return_value=diffusion_client, + ) + mock_remote_diff = mocker.patch.object(engine, "_create_remote_diffusion_stage") + mocker.patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(1)) + mocker.patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(1)) + mocker.patch.object( + engine, + "_attach_llm_stage", + return_value=(mocker.Mock(), mocker.Mock(), mocker.Mock(), mocker.Mock()), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.OmniMasterServer", + return_value=mock_oms, + ) + mocker.patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment") + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", + return_value=None, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", + return_value={}, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.extract_stage_metadata", + side_effect=lambda cfg: self._fake_metadata( + mocker, + cfg.stage_id, + getattr(cfg, "stage_type", "llm"), ), - patch( - "vllm_omni.engine.async_omni_engine.extract_stage_metadata", - side_effect=lambda cfg: self._fake_metadata(cfg.stage_id, getattr(cfg, "stage_type", "llm")), - ), - patch("vllm_omni.engine.async_omni_engine.finalize_initialized_stages", return_value=finalized), - ): - engine._initialize_stages(stage_init_timeout=60) + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", + return_value=finalized, + ) + + engine._initialize_stages(stage_init_timeout=60) assert mock_local_diff.call_count == 1 assert mock_local_diff.call_args.args[1].stage_id == 0 mock_remote_diff.assert_not_called() - def test_non_matching_diffusion_stage_uses_remote_diffusion_client(self): + def test_non_matching_diffusion_stage_uses_remote_diffusion_client(self, mocker: MockerFixture): """A non-local diffusion stage should attach via the remote diffusion path.""" stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1, stage_type="diffusion")] engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=0) - mock_oms = Mock(spec=OmniMasterServer) - remote_diffusion_client = Mock(stage_type="diffusion") + mock_oms = mocker.Mock(spec=OmniMasterServer) + remote_diffusion_client = mocker.Mock(stage_type="diffusion") finalized = ( - [Mock(), remote_diffusion_client], - [Mock(), Mock()], + [mocker.Mock(), remote_diffusion_client], + [mocker.Mock(), mocker.Mock()], [ {"final_output": False, "final_output_type": None, "stage_type": "llm"}, {"final_output": False, "final_output_type": None, "stage_type": "diffusion"}, ], ) - with ( - patch.object(engine, "_launch_diffusion_stage") as mock_local_diff, - patch.object( - engine, "_create_remote_diffusion_stage", return_value=remote_diffusion_client - ) as mock_remote_diff, - patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(0)), - patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(0)), - patch.object(engine, "_attach_llm_stage", return_value=(Mock(), Mock(), Mock(), Mock())), - patch("vllm_omni.engine.async_omni_engine.OmniMasterServer", return_value=mock_oms), - patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment"), - patch("vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", return_value=None), - patch("vllm_omni.engine.async_omni_engine.get_stage_connector_spec", return_value={}), - patch( - "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", return_value=(None, None, None) + mock_local_diff = mocker.patch.object(engine, "_launch_diffusion_stage") + mock_remote_diff = mocker.patch.object( + engine, + "_create_remote_diffusion_stage", + return_value=remote_diffusion_client, + ) + mocker.patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(0)) + mocker.patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(0)) + mocker.patch.object( + engine, + "_attach_llm_stage", + return_value=(mocker.Mock(), mocker.Mock(), mocker.Mock(), mocker.Mock()), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.OmniMasterServer", + return_value=mock_oms, + ) + mocker.patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment") + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", + return_value=None, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", + return_value={}, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.extract_stage_metadata", + side_effect=lambda cfg: self._fake_metadata( + mocker, + cfg.stage_id, + getattr(cfg, "stage_type", "llm"), ), - patch( - "vllm_omni.engine.async_omni_engine.extract_stage_metadata", - side_effect=lambda cfg: self._fake_metadata(cfg.stage_id, getattr(cfg, "stage_type", "llm")), - ), - patch("vllm_omni.engine.async_omni_engine.finalize_initialized_stages", return_value=finalized), - ): - engine._initialize_stages(stage_init_timeout=60) + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", + return_value=finalized, + ) + + engine._initialize_stages(stage_init_timeout=60) mock_local_diff.assert_not_called() assert mock_remote_diff.call_count == 1 @@ -894,45 +1069,47 @@ class TestInitializeStagesRouting: class TestLaunchDiffusionStage: """Test local diffusion stage launch wiring.""" - def test_registers_stage_with_public_master_properties(self): + def test_registers_stage_with_public_master_properties(self, mocker: MockerFixture): engine = object.__new__(AsyncOmniEngine) engine.model = "fake-model" engine.diffusion_batch_size = 4 stage_cfg = _make_stage_cfg(5, stage_type="diffusion") - metadata = Mock(stage_id=5) - omni_master_server = Mock(spec=OmniMasterServer) + metadata = mocker.Mock(stage_id=5) + omni_master_server = mocker.Mock(spec=OmniMasterServer) omni_master_server.address = "127.0.0.1" omni_master_server.port = 25000 - proc = Mock() - diffusion_client = Mock() + proc = mocker.Mock() + diffusion_client = mocker.Mock() - with ( - patch("vllm_omni.engine.async_omni_engine.build_diffusion_config", return_value="diffusion-config"), - patch( - "vllm_omni.engine.async_omni_engine.register_stage_with_omni_master", - return_value=( - "tcp://127.0.0.1:25001", - "tcp://127.0.0.1:25002", - "tcp://127.0.0.1:25003", - ), - ) as mock_register, - patch( - "vllm_omni.engine.async_omni_engine.spawn_diffusion_proc", - return_value=(proc, None, None, None), - ) as mock_spawn, - patch("vllm_omni.engine.async_omni_engine.complete_diffusion_handshake") as mock_handshake, - patch( - "vllm_omni.engine.async_omni_engine.StageDiffusionClient.from_addresses", - return_value=diffusion_client, - ) as mock_from_addresses, - ): - result = engine._launch_diffusion_stage( - stage_cfg=stage_cfg, - metadata=metadata, - omni_master_server=omni_master_server, - ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.build_diffusion_config", + return_value="diffusion-config", + ) + mock_register = mocker.patch( + "vllm_omni.engine.async_omni_engine.register_stage_with_omni_master", + return_value=( + "tcp://127.0.0.1:25001", + "tcp://127.0.0.1:25002", + "tcp://127.0.0.1:25003", + ), + ) + mock_spawn = mocker.patch( + "vllm_omni.engine.async_omni_engine.spawn_diffusion_proc", + return_value=(proc, None, None, None), + ) + mock_handshake = mocker.patch("vllm_omni.engine.async_omni_engine.complete_diffusion_handshake") + mock_from_addresses = mocker.patch( + "vllm_omni.engine.async_omni_engine.StageDiffusionClient.from_addresses", + return_value=diffusion_client, + ) + + result = engine._launch_diffusion_stage( + stage_cfg=stage_cfg, + metadata=metadata, + omni_master_server=omni_master_server, + ) mock_register.assert_called_once_with( omni_master_address="127.0.0.1", @@ -967,14 +1144,14 @@ class TestLaunchDiffusionStage: class TestCreateRemoteLlmStage: """Test _create_remote_llm_stage delegates correctly.""" - def _engine(self) -> AsyncOmniEngine: + def _engine(self, mocker: MockerFixture) -> AsyncOmniEngine: engine = object.__new__(AsyncOmniEngine) engine.model = "fake-model" engine.single_stage_mode = True engine._single_stage_id_filter = 0 - engine._omni_master_server = Mock(spec=OmniMasterServer) - engine._omni_master_server.get_zmq_addresses.return_value = Mock() - engine._omni_master_server.get_allocation.return_value = Mock() + engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) + engine._omni_master_server.get_zmq_addresses.return_value = mocker.Mock() + engine._omni_master_server.get_allocation.return_value = mocker.Mock() engine._omni_master_server.get_stage_config.return_value = { "stage_id": 0, "stage_type": "llm", @@ -982,42 +1159,40 @@ class TestCreateRemoteLlmStage: } return engine - @contextmanager - def _patch_build_and_connect(self, stage_id: int): - fake_vllm_config = Mock() - fake_executor_cls = Mock() - fake_addresses = Mock() + def _mock_build_and_connect(self, mocker: MockerFixture, stage_id: int): + fake_vllm_config = mocker.Mock() + fake_executor_cls = mocker.Mock() + fake_addresses = mocker.Mock() fake_addresses.inputs = ["tcp://127.0.0.1:5000"] fake_addresses.outputs = ["tcp://127.0.0.1:5001"] fake_addresses.frontend_stats_publish_address = None - eng_mgr = Mock() - coordinator = Mock() + eng_mgr = mocker.Mock() + coordinator = mocker.Mock() @contextmanager def fake_connect_cm(*args, **kwargs): yield eng_mgr, coordinator, fake_addresses - with ( - patch( - "vllm_omni.engine.async_omni_engine.build_engine_args_dict", - return_value={"model": "fake", "stage_id": stage_id}, - ), - patch( - "vllm_omni.engine.async_omni_engine.build_vllm_config", - return_value=(fake_vllm_config, fake_executor_cls), - ), - patch( - "vllm_omni.engine.async_omni_engine.connect_remote_engine_cores", - return_value=fake_connect_cm(), - ) as mock_connect, - ): - yield mock_connect, fake_vllm_config, fake_executor_cls, fake_addresses + mocker.patch( + "vllm_omni.engine.async_omni_engine.build_engine_args_dict", + return_value={"model": "fake", "stage_id": stage_id}, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.build_vllm_config", + return_value=(fake_vllm_config, fake_executor_cls), + ) + mock_connect = mocker.patch( + "vllm_omni.engine.async_omni_engine.connect_remote_engine_cores", + return_value=fake_connect_cm(), + ) - def test_returns_started_llm_stage_with_correct_stage_id(self): - engine = self._engine() + return mock_connect, fake_vllm_config, fake_executor_cls, fake_addresses + + def test_returns_started_llm_stage_with_correct_stage_id(self, mocker: MockerFixture): + engine = self._engine(mocker) stage_cfg = _make_stage_cfg(1) - metadata = Mock(stage_id=1) + metadata = mocker.Mock(stage_id=1) omni_ms = engine._omni_master_server omni_ms.get_stage_config.return_value = { "stage_id": 1, @@ -1025,52 +1200,77 @@ class TestCreateRemoteLlmStage: "engine_args": {}, } - with self._patch_build_and_connect(1): - result = engine._create_remote_llm_stage( - stage_cfg=stage_cfg, - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=60, - omni_master_server=omni_ms, - ) + self._mock_build_and_connect(mocker, 1) + result = engine._create_remote_llm_stage( + stage_cfg=stage_cfg, + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=60, + omni_master_server=omni_ms, + ) assert isinstance(result, StartedLlmStage) assert result.stage_id == 1 - def test_connect_remote_engine_cores_called_with_stage_id(self): - engine = self._engine() + def test_connect_remote_engine_cores_called_with_stage_id(self, mocker: MockerFixture): + engine = self._engine(mocker) stage_cfg = _make_stage_cfg(2) - metadata = Mock(stage_id=2) + metadata = mocker.Mock(stage_id=2) omni_ms = engine._omni_master_server - omni_ms.get_zmq_addresses.return_value = Mock(inputs=["x"], outputs=["y"]) + omni_ms.get_zmq_addresses.return_value = mocker.Mock(inputs=["x"], outputs=["y"]) omni_ms.get_stage_config.return_value = { "stage_id": 2, "stage_type": "llm", "engine_args": {}, } - fake_vllm_config = Mock() - fake_executor_cls = Mock() - fake_addresses = Mock() + fake_vllm_config = mocker.Mock() + fake_executor_cls = mocker.Mock() + fake_addresses = mocker.Mock() fake_addresses.inputs = ["tcp://127.0.0.1:5000"] fake_addresses.outputs = ["tcp://127.0.0.1:5001"] fake_addresses.frontend_stats_publish_address = None @contextmanager def fake_connect_cm(*args, **kwargs): - yield Mock(), Mock(), fake_addresses + yield mocker.Mock(), mocker.Mock(), fake_addresses - with ( - patch( - "vllm_omni.engine.async_omni_engine.build_engine_args_dict", - return_value={"model": "fake", "stage_id": 2}, - ), - patch( - "vllm_omni.engine.async_omni_engine.build_vllm_config", - return_value=(fake_vllm_config, fake_executor_cls), - ), - patch( - "vllm_omni.engine.async_omni_engine.connect_remote_engine_cores", return_value=fake_connect_cm() - ) as mock_connect, + mocker.patch( + "vllm_omni.engine.async_omni_engine.build_engine_args_dict", + return_value={"model": "fake", "stage_id": 2}, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.build_vllm_config", + return_value=(fake_vllm_config, fake_executor_cls), + ) + mock_connect = mocker.patch( + "vllm_omni.engine.async_omni_engine.connect_remote_engine_cores", + return_value=fake_connect_cm(), + ) + + engine._create_remote_llm_stage( + stage_cfg=stage_cfg, + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=60, + omni_master_server=omni_ms, + ) + + mock_connect.assert_called_once() + _, kwargs = mock_connect.call_args + assert kwargs.get("stage_id") == 2 or mock_connect.call_args.args[-1] == 2 + omni_ms.get_stage_config.assert_called_once_with(2, timeout_s=60) + + def test_missing_registered_stage_config_raises_value_error(self, mocker: MockerFixture): + engine = self._engine(mocker) + stage_cfg = _make_stage_cfg(3) + metadata = mocker.Mock(stage_id=3) + omni_ms = engine._omni_master_server + omni_ms.get_stage_config.return_value = None + + mock_build_args = mocker.patch("vllm_omni.engine.async_omni_engine.build_engine_args_dict") + with pytest.raises( + ValueError, + match="Remote stage 3 registered without stage config", ): engine._create_remote_llm_stage( stage_cfg=stage_cfg, @@ -1080,38 +1280,13 @@ class TestCreateRemoteLlmStage: omni_master_server=omni_ms, ) - mock_connect.assert_called_once() - _, kwargs = mock_connect.call_args - assert kwargs.get("stage_id") == 2 or mock_connect.call_args.args[-1] == 2 - omni_ms.get_stage_config.assert_called_once_with(2, timeout_s=60) - - def test_missing_registered_stage_config_raises_value_error(self): - engine = self._engine() - stage_cfg = _make_stage_cfg(3) - metadata = Mock(stage_id=3) - omni_ms = engine._omni_master_server - omni_ms.get_stage_config.return_value = None - - with patch("vllm_omni.engine.async_omni_engine.build_engine_args_dict") as mock_build_args: - with pytest.raises( - ValueError, - match="Remote stage 3 registered without stage config", - ): - engine._create_remote_llm_stage( - stage_cfg=stage_cfg, - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=60, - omni_master_server=omni_ms, - ) - mock_build_args.assert_not_called() - def test_exception_during_connect_closes_started_stage(self): + def test_exception_during_connect_closes_started_stage(self, mocker: MockerFixture): """If an error occurs after StartedLlmStage creation, close_started_llm_stage is called.""" - engine = self._engine() + engine = self._engine(mocker) stage_cfg = _make_stage_cfg(1) - metadata = Mock(stage_id=1) + metadata = mocker.Mock(stage_id=1) omni_ms = engine._omni_master_server omni_ms.get_stage_config.return_value = { "stage_id": 1, @@ -1121,26 +1296,30 @@ class TestCreateRemoteLlmStage: @contextmanager def boom(*args, **kwargs): - yield Mock(), Mock(), Mock() + yield mocker.Mock(), mocker.Mock(), mocker.Mock() raise RuntimeError("handshake failed") - with ( - patch( - "vllm_omni.engine.async_omni_engine.build_engine_args_dict", - return_value={"model": "fake", "stage_id": 1}, - ), - patch("vllm_omni.engine.async_omni_engine.build_vllm_config", return_value=(Mock(), Mock())), - patch("vllm_omni.engine.async_omni_engine.connect_remote_engine_cores", return_value=boom()), - patch("vllm_omni.engine.async_omni_engine.close_started_llm_stage") as mock_close, - ): - with pytest.raises(RuntimeError, match="handshake failed"): - engine._create_remote_llm_stage( - stage_cfg=stage_cfg, - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=60, - omni_master_server=omni_ms, - ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.build_engine_args_dict", + return_value={"model": "fake", "stage_id": 1}, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.build_vllm_config", + return_value=(mocker.Mock(), mocker.Mock()), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.connect_remote_engine_cores", + return_value=boom(), + ) + mock_close = mocker.patch("vllm_omni.engine.async_omni_engine.close_started_llm_stage") + with pytest.raises(RuntimeError, match="handshake failed"): + engine._create_remote_llm_stage( + stage_cfg=stage_cfg, + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=60, + omni_master_server=omni_ms, + ) mock_close.assert_called_once() @@ -1148,27 +1327,29 @@ class TestConnectRemoteEngineCoresCoordinator: """Test coordinator launch parity with launch_core_engines.""" @staticmethod - def _build_vllm_config(*, dp_rank: int = 0, offline_mode: bool = False, needs_dp_coordinator: bool = True) -> Mock: - parallel_config = Mock() + def _build_vllm_config( + mocker: MockerFixture, *, dp_rank: int = 0, offline_mode: bool = False, needs_dp_coordinator: bool = True + ) -> Any: + parallel_config = mocker.Mock() parallel_config.data_parallel_size_local = 1 parallel_config.data_parallel_size = 2 parallel_config.data_parallel_rank = dp_rank parallel_config.data_parallel_rank_local = 0 if offline_mode else None - vllm_config = Mock() + vllm_config = mocker.Mock() vllm_config.parallel_config = parallel_config vllm_config.needs_dp_coordinator = needs_dp_coordinator - vllm_config.model_config = Mock(is_moe=False) + vllm_config.model_config = mocker.Mock(is_moe=False) return vllm_config - def test_uses_registered_coordinator_addresses(self): - vllm_config = self._build_vllm_config(dp_rank=0, offline_mode=False, needs_dp_coordinator=True) + def test_uses_registered_coordinator_addresses(self, mocker: MockerFixture): + vllm_config = self._build_vllm_config(mocker, dp_rank=0, offline_mode=False, needs_dp_coordinator=True) - omni_master_server = Mock(spec=OmniMasterServer) + omni_master_server = mocker.Mock(spec=OmniMasterServer) omni_master_server.get_zmq_addresses.return_value = EngineZmqAddresses( inputs=["tcp://client-in"], outputs=["tcp://client-out"] ) - omni_master_server.get_allocation.return_value = Mock(handshake_bind_address="tcp://127.0.0.1:26001") + omni_master_server.get_allocation.return_value = mocker.Mock(handshake_bind_address="tcp://127.0.0.1:26001") omni_master_server.get_stage_coordinator_addresses.return_value = StageCoordinatorAddresses( coordinator_input="tcp://coord-in", coordinator_output="tcp://coord-out", @@ -1177,103 +1358,107 @@ class TestConnectRemoteEngineCoresCoordinator: @contextmanager def fake_socket_ctx(*args, **kwargs): - yield Mock() + yield mocker.Mock() - with ( - patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()), - patch("vllm_omni.engine.stage_engine_startup._wait_for_omni_engine_startup") as mock_wait, - ): - with connect_remote_engine_cores( - vllm_config=vllm_config, - omni_master_server=omni_master_server, - stage_id=7, - ) as (_, yielded_coordinator, yielded_addresses): - assert yielded_coordinator is None - assert yielded_addresses.coordinator_input == "tcp://coord-in" - assert yielded_addresses.coordinator_output == "tcp://coord-out" - assert yielded_addresses.frontend_stats_publish_address == "tcp://stats" + mocker.patch( + "vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", + return_value=fake_socket_ctx(), + ) + mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup._wait_for_omni_engine_startup") + with connect_remote_engine_cores( + vllm_config=vllm_config, + omni_master_server=omni_master_server, + stage_id=7, + ) as (_, yielded_coordinator, yielded_addresses): + assert yielded_coordinator is None + assert yielded_addresses.coordinator_input == "tcp://coord-in" + assert yielded_addresses.coordinator_output == "tcp://coord-out" + assert yielded_addresses.frontend_stats_publish_address == "tcp://stats" omni_master_server.get_stage_coordinator_addresses.assert_called_once_with(7) mock_wait.assert_called_once() - def test_defaults_to_no_coordinator_addresses_when_none_registered(self): + def test_defaults_to_no_coordinator_addresses_when_none_registered(self, mocker: MockerFixture): vllm_config = self._build_vllm_config( + mocker, dp_rank=0, offline_mode=False, needs_dp_coordinator=True, ) - omni_master_server = Mock(spec=OmniMasterServer) + omni_master_server = mocker.Mock(spec=OmniMasterServer) omni_master_server.get_zmq_addresses.return_value = EngineZmqAddresses( inputs=["tcp://client-in"], outputs=["tcp://client-out"] ) - omni_master_server.get_allocation.return_value = Mock(handshake_bind_address="tcp://127.0.0.1:26001") + omni_master_server.get_allocation.return_value = mocker.Mock(handshake_bind_address="tcp://127.0.0.1:26001") omni_master_server.get_stage_coordinator_addresses.return_value = StageCoordinatorAddresses() @contextmanager def fake_socket_ctx(*args, **kwargs): - yield Mock() + yield mocker.Mock() - with ( - patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()), - patch("vllm_omni.engine.stage_engine_startup._wait_for_omni_engine_startup"), - ): - with connect_remote_engine_cores( - vllm_config=vllm_config, - omni_master_server=omni_master_server, - stage_id=7, - ) as (_, yielded_coordinator, yielded_addresses): - assert yielded_coordinator is None - assert yielded_addresses.coordinator_input is None - assert yielded_addresses.coordinator_output is None - assert yielded_addresses.frontend_stats_publish_address is None + mocker.patch( + "vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", + return_value=fake_socket_ctx(), + ) + mocker.patch("vllm_omni.engine.stage_engine_startup._wait_for_omni_engine_startup") + with connect_remote_engine_cores( + vllm_config=vllm_config, + omni_master_server=omni_master_server, + stage_id=7, + ) as (_, yielded_coordinator, yielded_addresses): + assert yielded_coordinator is None + assert yielded_addresses.coordinator_input is None + assert yielded_addresses.coordinator_output is None + assert yielded_addresses.frontend_stats_publish_address is None class TestLaunchOmniCoreEngines: """Tests for local omni engine launch wiring.""" - def test_registers_stage_once_and_reuses_handshake_for_all_local_engines(self): - parallel_config = Mock( + def test_registers_stage_once_and_reuses_handshake_for_all_local_engines(self, mocker: MockerFixture): + parallel_config = mocker.Mock( data_parallel_size_local=2, data_parallel_size=4, data_parallel_rank=3, ) - vllm_config = Mock(parallel_config=parallel_config) + vllm_config = mocker.Mock(parallel_config=parallel_config) - omni_master_server = Mock(spec=OmniMasterServer) + omni_master_server = mocker.Mock(spec=OmniMasterServer) omni_master_server.address = "127.0.0.1" omni_master_server.port = 26000 - omni_master_server.get_allocation.return_value = Mock(handshake_bind_address="tcp://127.0.0.1:26001") + omni_master_server.get_allocation.return_value = mocker.Mock(handshake_bind_address="tcp://127.0.0.1:26001") stage_config = {"stage_id": 7, "stage_type": "llm"} - local_engine_manager = Mock() + local_engine_manager = mocker.Mock() @contextmanager def fake_socket_ctx(*args, **kwargs): - yield Mock() + yield mocker.Mock() - with ( - patch( - "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", - return_value="tcp://127.0.0.1:26001", - ) as mock_register, - patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()), - patch( - "vllm_omni.engine.stage_engine_startup.CoreEngineProcManager", - return_value=local_engine_manager, - ) as mock_manager_cls, - patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup"), - ): - with launch_omni_core_engines( - vllm_config=vllm_config, - executor_class=Mock(), - log_stats=False, - omni_master_server=omni_master_server, - stage_id=7, - stage_config=stage_config, - ) as (yielded_manager, yielded_coordinator, yielded_addresses): - assert yielded_manager is local_engine_manager - assert yielded_coordinator is None + mock_register = mocker.patch( + "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", + return_value="tcp://127.0.0.1:26001", + ) + mocker.patch( + "vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", + return_value=fake_socket_ctx(), + ) + mock_manager_cls = mocker.patch( + "vllm_omni.engine.stage_engine_startup.CoreEngineProcManager", + return_value=local_engine_manager, + ) + mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") + with launch_omni_core_engines( + vllm_config=vllm_config, + executor_class=mocker.Mock(), + log_stats=False, + omni_master_server=omni_master_server, + stage_id=7, + stage_config=stage_config, + ) as (yielded_manager, yielded_coordinator, yielded_addresses): + assert yielded_manager is local_engine_manager + assert yielded_coordinator is None mock_register.assert_called_once_with( omni_master_address="127.0.0.1", @@ -1292,55 +1477,56 @@ class TestLaunchOmniCoreEngines: assert manager_kwargs["handshake_address"] == "tcp://127.0.0.1:26001" assert manager_kwargs["executor_class"] is not None - def test_registers_stage_with_coordinator_when_started(self): - parallel_config = Mock( + def test_registers_stage_with_coordinator_when_started(self, mocker: MockerFixture): + parallel_config = mocker.Mock( data_parallel_size_local=1, data_parallel_size=2, data_parallel_rank=0, ) - vllm_config = Mock(parallel_config=parallel_config) + vllm_config = mocker.Mock(parallel_config=parallel_config) vllm_config.needs_dp_coordinator = True - vllm_config.model_config = Mock(is_moe=False) + vllm_config.model_config = mocker.Mock(is_moe=False) - omni_master_server = Mock(spec=OmniMasterServer) + omni_master_server = mocker.Mock(spec=OmniMasterServer) omni_master_server.address = "127.0.0.1" omni_master_server.port = 26000 omni_master_server.get_zmq_addresses.return_value = EngineZmqAddresses( inputs=["tcp://client-in"], outputs=["tcp://client-out"] ) - omni_master_server.get_allocation.return_value = Mock(handshake_bind_address="tcp://127.0.0.1:26001") + omni_master_server.get_allocation.return_value = mocker.Mock(handshake_bind_address="tcp://127.0.0.1:26001") - coordinator = Mock() + coordinator = mocker.Mock() coordinator.proc.pid = 1234 coordinator.get_engine_socket_addresses.return_value = ("tcp://coord-in", "tcp://coord-out") coordinator.get_stats_publish_address.return_value = "tcp://stats" @contextmanager def fake_socket_ctx(*args, **kwargs): - yield Mock() + yield mocker.Mock() - with ( - patch("vllm_omni.engine.stage_engine_startup.DPCoordinator", return_value=coordinator), - patch( - "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", - return_value="tcp://127.0.0.1:26001", - ) as mock_register, - patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()), - patch( - "vllm_omni.engine.stage_engine_startup.CoreEngineProcManager", - return_value=Mock(), - ) as mock_manager_cls, - patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") as mock_wait, + mocker.patch("vllm_omni.engine.stage_engine_startup.DPCoordinator", return_value=coordinator) + mock_register = mocker.patch( + "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", + return_value="tcp://127.0.0.1:26001", + ) + mocker.patch( + "vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", + return_value=fake_socket_ctx(), + ) + mock_manager_cls = mocker.patch( + "vllm_omni.engine.stage_engine_startup.CoreEngineProcManager", + return_value=mocker.Mock(), + ) + mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") + with launch_omni_core_engines( + vllm_config=vllm_config, + executor_class=mocker.Mock(), + log_stats=False, + omni_master_server=omni_master_server, + stage_id=7, + stage_config={"stage_id": 7}, ): - with launch_omni_core_engines( - vllm_config=vllm_config, - executor_class=Mock(), - log_stats=False, - omni_master_server=omni_master_server, - stage_id=7, - stage_config={"stage_id": 7}, - ): - pass + pass mock_register.assert_called_once_with( omni_master_address="127.0.0.1", @@ -1363,19 +1549,19 @@ class TestLaunchLlmStageSingleStageMode: """Test that _launch_llm_stage selects launch_omni_core_engines when single_stage_mode=True and _omni_master_server is set.""" - def _build_engine_with_oms(self) -> AsyncOmniEngine: + def _build_engine_with_oms(self, mocker: MockerFixture) -> AsyncOmniEngine: engine = object.__new__(AsyncOmniEngine) engine.model = "fake-model" engine.single_stage_mode = True engine._single_stage_id_filter = 0 engine._llm_stage_launch_lock = threading.Lock() - mock_oms = Mock(spec=OmniMasterServer) + mock_oms = mocker.Mock(spec=OmniMasterServer) mock_oms.address = "127.0.0.1" mock_oms.port = 25000 - alloc = Mock() + alloc = mocker.Mock() alloc.handshake_bind_address = "tcp://127.0.0.1:25001" mock_oms.get_allocation.return_value = alloc - fake_addresses = Mock() + fake_addresses = mocker.Mock() fake_addresses.inputs = ["tcp://127.0.0.1:5000"] fake_addresses.outputs = ["tcp://127.0.0.1:5001"] fake_addresses.frontend_stats_publish_address = None @@ -1383,66 +1569,60 @@ class TestLaunchLlmStageSingleStageMode: engine._omni_master_server = mock_oms return engine - @contextmanager - def _patch_launch_omni_cm(self, stage_id: int): - fake_vllm_config = Mock() - fake_executor_cls = Mock() - fake_addresses = Mock() + def _mock_launch_omni(self, mocker: MockerFixture, stage_id: int): + fake_vllm_config = mocker.Mock() + fake_executor_cls = mocker.Mock() + fake_addresses = mocker.Mock() fake_addresses.inputs = ["tcp://127.0.0.1:5000"] fake_addresses.outputs = ["tcp://127.0.0.1:5001"] fake_addresses.frontend_stats_publish_address = None - eng_mgr = Mock() + eng_mgr = mocker.Mock() @contextmanager def fake_launch_omni(*args, **kwargs): yield eng_mgr, None, fake_addresses - with ( - patch("vllm_omni.engine.async_omni_engine.setup_stage_devices"), - patch( - "vllm_omni.engine.async_omni_engine.build_engine_args_dict", - return_value={"model": "fake", "stage_id": stage_id}, - ), - patch( - "vllm_omni.engine.async_omni_engine.build_vllm_config", - return_value=(fake_vllm_config, fake_executor_cls), - ), - patch( - "vllm_omni.engine.async_omni_engine.acquire_device_locks", - return_value=[], - ), - patch( - "vllm_omni.engine.async_omni_engine.release_device_locks", - ), - patch( - "vllm_omni.engine.async_omni_engine.launch_omni_core_engines", - return_value=fake_launch_omni(), - ) as mock_launch_omni, - ): - yield mock_launch_omni + mocker.patch("vllm_omni.engine.async_omni_engine.setup_stage_devices") + mocker.patch( + "vllm_omni.engine.async_omni_engine.build_engine_args_dict", + return_value={"model": "fake", "stage_id": stage_id}, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.build_vllm_config", + return_value=(fake_vllm_config, fake_executor_cls), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.acquire_device_locks", + return_value=[], + ) + mocker.patch("vllm_omni.engine.async_omni_engine.release_device_locks") + return mocker.patch( + "vllm_omni.engine.async_omni_engine.launch_omni_core_engines", + return_value=fake_launch_omni(), + ) - def test_launch_omni_core_engines_used_in_single_stage_mode(self): + def test_launch_omni_core_engines_used_in_single_stage_mode(self, mocker: MockerFixture): """single_stage_mode + _omni_master_server → launch_omni_core_engines.""" - engine = self._build_engine_with_oms() - metadata = Mock(stage_id=0, runtime_cfg={}) + engine = self._build_engine_with_oms(mocker) + metadata = mocker.Mock(stage_id=0, runtime_cfg={}) stage_cfg = _make_stage_cfg(0) - with self._patch_launch_omni_cm(0) as mock_launch_omni: - result = engine._launch_llm_stage( - stage_cfg=stage_cfg, - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=60, - llm_stage_launch_lock=threading.Lock(), - ) + mock_launch_omni = self._mock_launch_omni(mocker, 0) + result = engine._launch_llm_stage( + stage_cfg=stage_cfg, + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=60, + llm_stage_launch_lock=threading.Lock(), + ) mock_launch_omni.assert_called_once() assert mock_launch_omni.call_args.kwargs["stage_config"] is stage_cfg assert isinstance(result, StartedLlmStage) assert result.stage_id == 0 - def test_spawn_stage_core_used_in_normal_mode(self): + def test_spawn_stage_core_used_in_normal_mode(self, mocker: MockerFixture): """~single_stage_mode → spawn_stage_core + complete_stage_handshake.""" engine = object.__new__(AsyncOmniEngine) engine.model = "fake-model" @@ -1450,44 +1630,45 @@ class TestLaunchLlmStageSingleStageMode: engine._omni_master_server = None engine._llm_stage_launch_lock = threading.Lock() - fake_vllm_config = Mock() - fake_executor_cls = Mock() - fake_addresses = Mock() + fake_vllm_config = mocker.Mock() + fake_executor_cls = mocker.Mock() + fake_addresses = mocker.Mock() fake_addresses.inputs = ["tcp://127.0.0.1:5000"] fake_addresses.outputs = ["tcp://127.0.0.1:5001"] fake_addresses.frontend_stats_publish_address = None - fake_proc = Mock() + fake_proc = mocker.Mock() fake_handshake_address = "ipc:///tmp/fake-handshake" stage_init_timeout = 60 - with ( - patch("vllm_omni.engine.async_omni_engine.setup_stage_devices"), - patch( - "vllm_omni.engine.async_omni_engine.build_engine_args_dict", - return_value={"model": "fake", "stage_id": 0}, - ), - patch( - "vllm_omni.engine.async_omni_engine.build_vllm_config", - return_value=(fake_vllm_config, fake_executor_cls), - ), - patch("vllm_omni.engine.async_omni_engine.acquire_device_locks", return_value=[]), - patch("vllm_omni.engine.async_omni_engine.release_device_locks"), - patch( - "vllm_omni.engine.async_omni_engine.spawn_stage_core", - return_value=(fake_addresses, fake_proc, fake_handshake_address), - ) as mock_spawn, - patch("vllm_omni.engine.async_omni_engine.complete_stage_handshake") as mock_handshake, - patch("vllm_omni.engine.async_omni_engine.launch_omni_core_engines") as mock_omni, - ): - metadata = Mock(stage_id=0, runtime_cfg={}) - result = engine._launch_llm_stage( - stage_cfg=_make_stage_cfg(0), - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=stage_init_timeout, - llm_stage_launch_lock=threading.Lock(), - ) + mocker.patch("vllm_omni.engine.async_omni_engine.setup_stage_devices") + mocker.patch( + "vllm_omni.engine.async_omni_engine.build_engine_args_dict", + return_value={"model": "fake", "stage_id": 0}, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.build_vllm_config", + return_value=(fake_vllm_config, fake_executor_cls), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.acquire_device_locks", + return_value=[], + ) + mocker.patch("vllm_omni.engine.async_omni_engine.release_device_locks") + mock_spawn = mocker.patch( + "vllm_omni.engine.async_omni_engine.spawn_stage_core", + return_value=(fake_addresses, fake_proc, fake_handshake_address), + ) + mock_handshake = mocker.patch("vllm_omni.engine.async_omni_engine.complete_stage_handshake") + mock_omni = mocker.patch("vllm_omni.engine.async_omni_engine.launch_omni_core_engines") + metadata = mocker.Mock(stage_id=0, runtime_cfg={}) + result = engine._launch_llm_stage( + stage_cfg=_make_stage_cfg(0), + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=stage_init_timeout, + llm_stage_launch_lock=threading.Lock(), + ) mock_spawn.assert_called_once_with( vllm_config=fake_vllm_config, @@ -1505,50 +1686,58 @@ class TestLaunchLlmStageSingleStageMode: assert isinstance(result, StartedLlmStage) assert result.proc is fake_proc - def test_launch_omni_passes_stage_id_and_master_server(self): + def test_launch_omni_passes_stage_id_and_master_server(self, mocker: MockerFixture): """launch_omni_core_engines receives the correct stage_id and omni_master_server.""" - engine = self._build_engine_with_oms() - metadata = Mock(stage_id=0, runtime_cfg={}) + engine = self._build_engine_with_oms(mocker) + metadata = mocker.Mock(stage_id=0, runtime_cfg={}) captured_kwargs: dict[str, Any] = {} @contextmanager def capturing_launch(*args, **kwargs): captured_kwargs.update(kwargs) - fake_addresses = Mock() + fake_addresses = mocker.Mock() fake_addresses.inputs = ["tcp://127.0.0.1:5000"] fake_addresses.outputs = ["tcp://127.0.0.1:5001"] fake_addresses.frontend_stats_publish_address = None - yield Mock(), None, fake_addresses + yield mocker.Mock(), None, fake_addresses - with ( - patch("vllm_omni.engine.async_omni_engine.setup_stage_devices"), - patch( - "vllm_omni.engine.async_omni_engine.build_engine_args_dict", - return_value={"model": "fake", "stage_id": 0}, - ), - patch("vllm_omni.engine.async_omni_engine.build_vllm_config", return_value=(Mock(), Mock())), - patch("vllm_omni.engine.async_omni_engine.acquire_device_locks", return_value=[]), - patch("vllm_omni.engine.async_omni_engine.release_device_locks"), - patch("vllm_omni.engine.async_omni_engine.launch_omni_core_engines", side_effect=capturing_launch), - ): - engine._launch_llm_stage( - stage_cfg=_make_stage_cfg(0), - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=60, - llm_stage_launch_lock=threading.Lock(), - ) + mocker.patch("vllm_omni.engine.async_omni_engine.setup_stage_devices") + mocker.patch( + "vllm_omni.engine.async_omni_engine.build_engine_args_dict", + return_value={"model": "fake", "stage_id": 0}, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.build_vllm_config", + return_value=(mocker.Mock(), mocker.Mock()), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.acquire_device_locks", + return_value=[], + ) + mocker.patch("vllm_omni.engine.async_omni_engine.release_device_locks") + mocker.patch( + "vllm_omni.engine.async_omni_engine.launch_omni_core_engines", + side_effect=capturing_launch, + ) + + engine._launch_llm_stage( + stage_cfg=_make_stage_cfg(0), + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=60, + llm_stage_launch_lock=threading.Lock(), + ) assert captured_kwargs.get("stage_id") == 0 assert captured_kwargs.get("omni_master_server") is engine._omni_master_server - def test_launch_omni_context_exits_before_stage_cleanup_on_error(self): + def test_launch_omni_context_exits_before_stage_cleanup_on_error(self, mocker: MockerFixture): """Errors after entering the omni launch context still unwind it first.""" - engine = self._build_engine_with_oms() - metadata = Mock(stage_id=0, runtime_cfg={}) + engine = self._build_engine_with_oms(mocker) + metadata = mocker.Mock(stage_id=0, runtime_cfg={}) - fake_addresses = Mock() + fake_addresses = mocker.Mock() fake_addresses.inputs = ["tcp://127.0.0.1:5000"] fake_addresses.outputs = ["tcp://127.0.0.1:5001"] fake_addresses.frontend_stats_publish_address = None @@ -1558,47 +1747,51 @@ class TestLaunchLlmStageSingleStageMode: @contextmanager def fake_launch_omni(*args, **kwargs): try: - yield Mock(), None, fake_addresses + yield mocker.Mock(), None, fake_addresses finally: events.append("launch_exit") - with ( - patch("vllm_omni.engine.async_omni_engine.setup_stage_devices"), - patch( - "vllm_omni.engine.async_omni_engine.build_engine_args_dict", - return_value={"model": "fake", "stage_id": 0}, - ), - patch("vllm_omni.engine.async_omni_engine.build_vllm_config", return_value=(Mock(), Mock())), - patch("vllm_omni.engine.async_omni_engine.acquire_device_locks", return_value=[]), - patch("vllm_omni.engine.async_omni_engine.release_device_locks"), - patch( - "vllm_omni.engine.async_omni_engine.launch_omni_core_engines", - return_value=fake_launch_omni(), - ), - patch("vllm_omni.engine.async_omni_engine.logger.info", side_effect=RuntimeError("boom")), - patch( - "vllm_omni.engine.async_omni_engine.close_started_llm_stage", - side_effect=lambda _started: events.append("stage_close"), - ) as mock_close_stage, - ): - with pytest.raises(RuntimeError, match="boom"): - engine._launch_llm_stage( - stage_cfg=_make_stage_cfg(0), - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=60, - llm_stage_launch_lock=threading.Lock(), - ) + mocker.patch("vllm_omni.engine.async_omni_engine.setup_stage_devices") + mocker.patch( + "vllm_omni.engine.async_omni_engine.build_engine_args_dict", + return_value={"model": "fake", "stage_id": 0}, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.build_vllm_config", + return_value=(mocker.Mock(), mocker.Mock()), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.acquire_device_locks", + return_value=[], + ) + mocker.patch("vllm_omni.engine.async_omni_engine.release_device_locks") + mocker.patch( + "vllm_omni.engine.async_omni_engine.launch_omni_core_engines", + return_value=fake_launch_omni(), + ) + mocker.patch("vllm_omni.engine.async_omni_engine.logger.info", side_effect=RuntimeError("boom")) + mock_close_stage = mocker.patch( + "vllm_omni.engine.async_omni_engine.close_started_llm_stage", + side_effect=lambda _started: events.append("stage_close"), + ) + with pytest.raises(RuntimeError, match="boom"): + engine._launch_llm_stage( + stage_cfg=_make_stage_cfg(0), + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=60, + llm_stage_launch_lock=threading.Lock(), + ) mock_close_stage.assert_called_once() assert events == ["launch_exit", "stage_close"] - def test_base_exception_propagates_without_started_stage_cleanup(self): + def test_base_exception_propagates_without_started_stage_cleanup(self, mocker: MockerFixture): """BaseException subclasses should bypass the Exception cleanup path.""" - engine = self._build_engine_with_oms() - metadata = Mock(stage_id=0, runtime_cfg={}) + engine = self._build_engine_with_oms(mocker) + metadata = mocker.Mock(stage_id=0, runtime_cfg={}) - fake_addresses = Mock() + fake_addresses = mocker.Mock() fake_addresses.inputs = ["tcp://127.0.0.1:5000"] fake_addresses.outputs = ["tcp://127.0.0.1:5001"] fake_addresses.frontend_stats_publish_address = None @@ -1611,37 +1804,41 @@ class TestLaunchLlmStageSingleStageMode: @contextmanager def fake_launch_omni(*args, **kwargs): try: - yield Mock(), None, fake_addresses + yield mocker.Mock(), None, fake_addresses finally: events.append("launch_exit") - with ( - patch("vllm_omni.engine.async_omni_engine.setup_stage_devices"), - patch( - "vllm_omni.engine.async_omni_engine.build_engine_args_dict", - return_value={"model": "fake", "stage_id": 0}, - ), - patch("vllm_omni.engine.async_omni_engine.build_vllm_config", return_value=(Mock(), Mock())), - patch("vllm_omni.engine.async_omni_engine.acquire_device_locks", return_value=[]), - patch("vllm_omni.engine.async_omni_engine.release_device_locks"), - patch( - "vllm_omni.engine.async_omni_engine.launch_omni_core_engines", - return_value=fake_launch_omni(), - ), - patch( - "vllm_omni.engine.async_omni_engine.logger.info", - side_effect=FatalLaunchInterrupt("stop"), - ), - patch("vllm_omni.engine.async_omni_engine.close_started_llm_stage") as mock_close_stage, - ): - with pytest.raises(FatalLaunchInterrupt, match="stop"): - engine._launch_llm_stage( - stage_cfg=_make_stage_cfg(0), - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=60, - llm_stage_launch_lock=threading.Lock(), - ) + mocker.patch("vllm_omni.engine.async_omni_engine.setup_stage_devices") + mocker.patch( + "vllm_omni.engine.async_omni_engine.build_engine_args_dict", + return_value={"model": "fake", "stage_id": 0}, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.build_vllm_config", + return_value=(mocker.Mock(), mocker.Mock()), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.acquire_device_locks", + return_value=[], + ) + mocker.patch("vllm_omni.engine.async_omni_engine.release_device_locks") + mocker.patch( + "vllm_omni.engine.async_omni_engine.launch_omni_core_engines", + return_value=fake_launch_omni(), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.logger.info", + side_effect=FatalLaunchInterrupt("stop"), + ) + mock_close_stage = mocker.patch("vllm_omni.engine.async_omni_engine.close_started_llm_stage") + with pytest.raises(FatalLaunchInterrupt, match="stop"): + engine._launch_llm_stage( + stage_cfg=_make_stage_cfg(0), + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=60, + llm_stage_launch_lock=threading.Lock(), + ) mock_close_stage.assert_not_called() assert events == ["launch_exit"] diff --git a/tests/entrypoints/openai_api/test_serving_chat_speaker.py b/tests/entrypoints/openai_api/test_serving_chat_speaker.py index 3b9151120..97c05e45b 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_speaker.py +++ b/tests/entrypoints/openai_api/test_serving_chat_speaker.py @@ -4,9 +4,9 @@ import asyncio from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock import pytest +from pytest_mock import MockerFixture from vllm_omni.entrypoints.openai.utils import ( get_supported_speakers_from_hf_config, @@ -25,9 +25,9 @@ def serving_chat(): return instance -def _make_hf_config(*, speaker_id: dict | None = None, spk_id: dict | None = None): - hf_config = MagicMock() - talker_config = MagicMock() +def _make_hf_config(mocker: MockerFixture, *, speaker_id: dict | None = None, spk_id: dict | None = None): + hf_config = mocker.MagicMock() + talker_config = mocker.MagicMock() talker_config.speaker_id = speaker_id talker_config.spk_id = spk_id hf_config.talker_config = talker_config @@ -51,14 +51,14 @@ def test_validate_requested_speaker_skips_validation_when_supported_empty(): assert validate_requested_speaker(" ", {"vivian"}) is None -def test_get_supported_speakers_from_hf_config_uses_spk_id_fallback(): - hf_config = _make_hf_config(speaker_id=None, spk_id={"Serena": 0}) +def test_get_supported_speakers_from_hf_config_uses_spk_id_fallback(mocker: MockerFixture): + hf_config = _make_hf_config(mocker, speaker_id=None, spk_id={"Serena": 0}) assert get_supported_speakers_from_hf_config(hf_config) == {"serena"} -def test_get_supported_speakers_caches_normalized_keys(serving_chat): - serving_chat.model_config = MagicMock() - serving_chat.model_config.hf_config = _make_hf_config(speaker_id={"Vivian": 0, "Ethan": 1}) +def test_get_supported_speakers_caches_normalized_keys(mocker: MockerFixture, serving_chat): + serving_chat.model_config = mocker.MagicMock() + serving_chat.model_config.hf_config = _make_hf_config(mocker, speaker_id={"Vivian": 0, "Ethan": 1}) assert serving_chat._get_supported_speakers() == {"vivian", "ethan"} @@ -67,15 +67,15 @@ def test_get_supported_speakers_caches_normalized_keys(serving_chat): assert serving_chat._get_supported_speakers() == {"vivian", "ethan"} -def test_create_chat_completion_converts_value_error_to_error_response(serving_chat): +def test_create_chat_completion_converts_value_error_to_error_response(mocker: MockerFixture, serving_chat): serving_chat._diffusion_mode = False - serving_chat._check_model = AsyncMock(return_value=None) - serving_chat.engine_client = MagicMock(errored=False) - serving_chat._maybe_get_adapters = MagicMock(return_value=None) - serving_chat.models = MagicMock() + serving_chat._check_model = mocker.AsyncMock(return_value=None) + serving_chat.engine_client = mocker.MagicMock(errored=False) + serving_chat._maybe_get_adapters = mocker.MagicMock(return_value=None) + serving_chat.models = mocker.MagicMock() serving_chat.models.model_name.return_value = "test-model" - serving_chat.renderer = MagicMock() - serving_chat.renderer.get_tokenizer.return_value = MagicMock() + serving_chat.renderer = mocker.MagicMock() + serving_chat.renderer.get_tokenizer.return_value = mocker.MagicMock() serving_chat.reasoning_parser_cls = None serving_chat.tool_parser = None serving_chat.use_harmony = False @@ -85,12 +85,12 @@ def test_create_chat_completion_converts_value_error_to_error_response(serving_c serving_chat.chat_template = None serving_chat.chat_template_content_format = "string" serving_chat.default_chat_template_kwargs = {} - serving_chat._validate_chat_template = MagicMock(return_value=None) - serving_chat._prepare_extra_chat_template_kwargs = MagicMock(return_value={}) - serving_chat._preprocess_chat = AsyncMock( + serving_chat._validate_chat_template = mocker.MagicMock(return_value=None) + serving_chat._prepare_extra_chat_template_kwargs = mocker.MagicMock(return_value={}) + serving_chat._preprocess_chat = mocker.AsyncMock( side_effect=ValueError("Invalid speaker 'uncle_fu'. Supported: ethan, vivian") ) - serving_chat.create_error_response = MagicMock(return_value="error-response") + serving_chat.create_error_response = mocker.MagicMock(return_value="error-response") request = SimpleNamespace( tool_choice=None, diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 06b6f5c16..c88412062 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -6,7 +6,6 @@ import struct from inspect import Signature, signature from pathlib import Path from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock, patch import numpy as np import pytest @@ -901,7 +900,7 @@ class TestTTSMethods: # Verify speakers are normalized to lowercase assert server.supported_speakers == {"ryan", "vivian", "aiden"} - def test_build_tts_params_with_uploaded_voice(self, speech_server): + def test_build_tts_params_with_uploaded_voice(self, speech_server, mocker: MockerFixture): """Test _build_tts_params auto-sets ref_audio for uploaded voices (x_vector only).""" speech_server.uploaded_speakers = { "custom_voice": { @@ -914,18 +913,18 @@ class TestTTSMethods: } speech_server.supported_speakers = {"ryan", "vivian", "custom_voice"} - with patch.object(speech_server, "_get_uploaded_audio_data") as mock_get_audio: - mock_get_audio.return_value = "data:audio/wav;base64,ZmFrZWF1ZGlv" - req = OpenAICreateSpeechRequest(input="Hello", voice="custom_voice") - params = speech_server._build_tts_params(req) + mock_get_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data") + mock_get_audio.return_value = "data:audio/wav;base64,ZmFrZWF1ZGlv" + req = OpenAICreateSpeechRequest(input="Hello", voice="custom_voice") + params = speech_server._build_tts_params(req) - assert params["ref_audio"] == ["data:audio/wav;base64,ZmFrZWF1ZGlv"] - assert params["x_vector_only_mode"] == [True] - assert params["task_type"] == ["Base"] - assert params["voice_created_at"] == [1711234567.89] - assert "ref_text" not in params + assert params["ref_audio"] == ["data:audio/wav;base64,ZmFrZWF1ZGlv"] + assert params["x_vector_only_mode"] == [True] + assert params["task_type"] == ["Base"] + assert params["voice_created_at"] == [1711234567.89] + assert "ref_text" not in params - def test_build_tts_params_with_uploaded_voice_ref_text(self, speech_server): + def test_build_tts_params_with_uploaded_voice_ref_text(self, speech_server, mocker: MockerFixture): """Test _build_tts_params enables in-context cloning when ref_text is stored.""" speech_server.uploaded_speakers = { "custom_voice": { @@ -938,16 +937,16 @@ class TestTTSMethods: } speech_server.supported_speakers = {"ryan", "vivian", "custom_voice"} - with patch.object(speech_server, "_get_uploaded_audio_data") as mock_get_audio: - mock_get_audio.return_value = "data:audio/wav;base64,ZmFrZWF1ZGlv" - req = OpenAICreateSpeechRequest(input="Hello", voice="custom_voice") - params = speech_server._build_tts_params(req) + mock_get_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data") + mock_get_audio.return_value = "data:audio/wav;base64,ZmFrZWF1ZGlv" + req = OpenAICreateSpeechRequest(input="Hello", voice="custom_voice") + params = speech_server._build_tts_params(req) - assert params["ref_audio"] == ["data:audio/wav;base64,ZmFrZWF1ZGlv"] - assert params["x_vector_only_mode"] == [False] - assert params["task_type"] == ["Base"] - assert params["ref_text"] == ["Hello world transcript"] - assert params["voice_created_at"] == [1711234567.89] + assert params["ref_audio"] == ["data:audio/wav;base64,ZmFrZWF1ZGlv"] + assert params["x_vector_only_mode"] == [False] + assert params["task_type"] == ["Base"] + assert params["ref_text"] == ["Hello world transcript"] + assert params["voice_created_at"] == [1711234567.89] def test_build_tts_params_without_uploaded_voice(self, speech_server): """Test _build_tts_params does not auto-set ref_audio for non-uploaded voices.""" @@ -989,45 +988,43 @@ class TestTTSMethods: # x_vector_only_mode should not be set when explicit ref_audio is provided assert "x_vector_only_mode" not in params - def test_get_uploaded_audio_data(self, speech_server): + def test_get_uploaded_audio_data(self, speech_server, mocker: MockerFixture): """Test _get_uploaded_audio_data function.""" # Mock file operations - with ( - patch("builtins.open", create=True) as mock_open, - patch("base64.b64encode") as mock_b64encode, - patch("pathlib.Path.exists") as mock_exists, - ): - mock_exists.return_value = True - mock_b64encode.return_value = b"ZmFrZWF1ZGlv" + mock_open = mocker.patch("builtins.open", create=True) + mock_b64encode = mocker.patch("base64.b64encode") + mock_exists = mocker.patch("pathlib.Path.exists") + mock_exists.return_value = True + mock_b64encode.return_value = b"ZmFrZWF1ZGlv" - # Setup mock file - mock_file = MagicMock() - mock_file.read.return_value = b"fakeaudio" - mock_open.return_value.__enter__.return_value = mock_file + # Setup mock file + mock_file = mocker.MagicMock() + mock_file.read.return_value = b"fakeaudio" + mock_open.return_value.__enter__.return_value = mock_file - # Setup uploaded speaker - speech_server.uploaded_speakers = { - "test_voice": {"name": "test_voice", "file_path": "/tmp/test.wav", "mime_type": "audio/wav"} - } - result = speech_server._get_uploaded_audio_data("test_voice") + # Setup uploaded speaker + speech_server.uploaded_speakers = { + "test_voice": {"name": "test_voice", "file_path": "/tmp/test.wav", "mime_type": "audio/wav"} + } + result = speech_server._get_uploaded_audio_data("test_voice") - assert result == "data:audio/wav;base64,ZmFrZWF1ZGlv" - mock_open.assert_called_once_with(Path("/tmp/test.wav"), "rb") - mock_b64encode.assert_called_once_with(b"fakeaudio") + assert result == "data:audio/wav;base64,ZmFrZWF1ZGlv" + mock_open.assert_called_once_with(Path("/tmp/test.wav"), "rb") + mock_b64encode.assert_called_once_with(b"fakeaudio") - def test_get_uploaded_audio_data_missing_file(self, speech_server): + def test_get_uploaded_audio_data_missing_file(self, speech_server, mocker: MockerFixture): """Test _get_uploaded_audio_data when file is missing.""" - with patch("pathlib.Path.exists") as mock_exists: - mock_exists.return_value = False + mock_exists = mocker.patch("pathlib.Path.exists") + mock_exists.return_value = False - # Setup uploaded speaker - speech_server.uploaded_speakers = { - "test_voice": {"name": "test_voice", "file_path": "/tmp/test.wav", "mime_type": "audio/wav"} - } + # Setup uploaded speaker + speech_server.uploaded_speakers = { + "test_voice": {"name": "test_voice", "file_path": "/tmp/test.wav", "mime_type": "audio/wav"} + } - result = speech_server._get_uploaded_audio_data("test_voice") + result = speech_server._get_uploaded_audio_data("test_voice") - assert result is None + assert result is None def test_get_uploaded_audio_data_voice_not_found(self, speech_server): """Test _get_uploaded_audio_data when voice is not in uploaded_speakers.""" @@ -1049,7 +1046,7 @@ class TestTTSMethods: req = OpenAICreateSpeechRequest.model_validate({"input": "Hello", "voice": "custom_voice"}) assert req.voice == "custom_voice" - def test_speaker_alias_in_base_task_with_uploaded_voice(self, speech_server): + def test_speaker_alias_in_base_task_with_uploaded_voice(self, speech_server, mocker: MockerFixture): """Using 'speaker' key with an uploaded voice should work for Base task.""" speech_server.uploaded_speakers = { "utesf": { @@ -1061,13 +1058,13 @@ class TestTTSMethods: } req = OpenAICreateSpeechRequest.model_validate({"input": "Hello", "speaker": "UTESF", "task_type": "Base"}) assert req.voice == "UTESF" - with patch("pathlib.Path.exists", return_value=True): - result = speech_server._validate_qwen_tts_request(req) + mocker.patch("pathlib.Path.exists", return_value=True) + result = speech_server._validate_qwen_tts_request(req) assert result is None # ── uploaded voice with embedding ── - def test_build_tts_params_with_uploaded_voice_embedding(self, speech_server): + def test_build_tts_params_with_uploaded_voice_embedding(self, speech_server, mocker: MockerFixture): """Test _build_tts_params loads embedding for embedding-uploaded voices.""" speech_server.uploaded_speakers = { "emb_voice": { @@ -1083,20 +1080,20 @@ class TestTTSMethods: speech_server.supported_speakers = {"ryan", "vivian", "emb_voice"} fake_embedding = [0.1] * 1024 - with patch.object(speech_server, "_get_uploaded_speaker_embedding") as mock_get_emb: - mock_get_emb.return_value = fake_embedding - req = OpenAICreateSpeechRequest(input="Hello", voice="emb_voice") - params = speech_server._build_tts_params(req) + mock_get_emb = mocker.patch.object(speech_server, "_get_uploaded_speaker_embedding") + mock_get_emb.return_value = fake_embedding + req = OpenAICreateSpeechRequest(input="Hello", voice="emb_voice") + params = speech_server._build_tts_params(req) - assert "voice_clone_prompt" in params - assert params["voice_clone_prompt"][0]["ref_spk_embedding"] == fake_embedding - assert params["task_type"] == ["Base"] - assert params["x_vector_only_mode"] == [True] - assert "ref_audio" not in params + assert "voice_clone_prompt" in params + assert params["voice_clone_prompt"][0]["ref_spk_embedding"] == fake_embedding + assert params["task_type"] == ["Base"] + assert params["x_vector_only_mode"] == [True] + assert "ref_audio" not in params # ── regression: full flow from issue #1603 ── - def test_regression_1603_speaker_key_with_uploaded_audio_voice(self, speech_server): + def test_regression_1603_speaker_key_with_uploaded_audio_voice(self, speech_server, mocker: MockerFixture): """Regression test for #1603: upload audio voice, then invoke TTS with 'speaker' key. Verifies the full validate → build_params pipeline works end-to-end. @@ -1116,14 +1113,14 @@ class TestTTSMethods: assert req.voice == "UTESF" # Validation should pass (file exists) - with patch("pathlib.Path.exists", return_value=True): - err = speech_server._validate_qwen_tts_request(req) + mocker.patch("pathlib.Path.exists", return_value=True) + err = speech_server._validate_qwen_tts_request(req) assert err is None, f"Validation failed: {err}" # Build params should auto-set ref_audio from stored file - with patch.object(speech_server, "_get_uploaded_audio_data") as mock_audio: - mock_audio.return_value = "data:audio/wav;base64,ZmFrZQ==" - params = speech_server._build_tts_params(req) + mock_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data") + mock_audio.return_value = "data:audio/wav;base64,ZmFrZQ==" + params = speech_server._build_tts_params(req) assert params["task_type"] == ["Base"] assert params["ref_audio"] == ["data:audio/wav;base64,ZmFrZQ=="] @@ -1131,7 +1128,7 @@ class TestTTSMethods: assert params["x_vector_only_mode"] == [False] assert params["speaker"] == ["utesf"] - def test_regression_1603_speaker_key_with_uploaded_embedding_voice(self, speech_server): + def test_regression_1603_speaker_key_with_uploaded_embedding_voice(self, speech_server, mocker: MockerFixture): """Regression test for #1603: upload embedding voice, then invoke TTS with 'speaker' key. Verifies embedding-uploaded voices are loaded as voice_clone_prompt, not as audio. @@ -1154,15 +1151,15 @@ class TestTTSMethods: assert req.voice == "myvoice" # Validation should pass - with patch("pathlib.Path.exists", return_value=True): - err = speech_server._validate_qwen_tts_request(req) + mocker.patch("pathlib.Path.exists", return_value=True) + err = speech_server._validate_qwen_tts_request(req) assert err is None, f"Validation failed: {err}" # Build params should use embedding, NOT audio fake_emb = [0.1] * 1024 - with patch.object(speech_server, "_get_uploaded_speaker_embedding") as mock_emb: - mock_emb.return_value = fake_emb - params = speech_server._build_tts_params(req) + mock_emb = mocker.patch.object(speech_server, "_get_uploaded_speaker_embedding") + mock_emb.return_value = fake_emb + params = speech_server._build_tts_params(req) assert params["task_type"] == ["Base"] assert params["x_vector_only_mode"] == [True] @@ -1171,7 +1168,7 @@ class TestTTSMethods: # Must NOT have ref_audio — that would fail for safetensors files assert "ref_audio" not in params - def test_validate_rejects_embedding_voice_with_pending_cache(self, speech_server): + def test_validate_rejects_embedding_voice_with_pending_cache(self, speech_server, mocker: MockerFixture): """Validation should reject embedding voices whose cache is not yet ready.""" speech_server.uploaded_speakers = { "myvoice": { @@ -1184,12 +1181,12 @@ class TestTTSMethods: } } req = OpenAICreateSpeechRequest.model_validate({"input": "Hello", "speaker": "myvoice", "task_type": "Base"}) - with patch("pathlib.Path.exists", return_value=True): - err = speech_server._validate_qwen_tts_request(req) + mocker.patch("pathlib.Path.exists", return_value=True) + err = speech_server._validate_qwen_tts_request(req) assert err is not None assert "not yet ready" in err - def test_x_vector_only_mode_not_overwritten_for_uploaded_embedding(self, speech_server): + def test_x_vector_only_mode_not_overwritten_for_uploaded_embedding(self, speech_server, mocker: MockerFixture): """x_vector_only_mode set by uploaded embedding must not be overwritten by request field.""" speech_server.uploaded_speakers = { "emb_voice": { @@ -1203,11 +1200,11 @@ class TestTTSMethods: } } fake_emb = [0.1] * 1024 - with patch.object(speech_server, "_get_uploaded_speaker_embedding") as mock_emb: - mock_emb.return_value = fake_emb - # Client explicitly sends x_vector_only_mode=False, but embedding requires True - req = OpenAICreateSpeechRequest(input="Hello", voice="emb_voice", x_vector_only_mode=False) - params = speech_server._build_tts_params(req) + mock_emb = mocker.patch.object(speech_server, "_get_uploaded_speaker_embedding") + mock_emb.return_value = fake_emb + # Client explicitly sends x_vector_only_mode=False, but embedding requires True + req = OpenAICreateSpeechRequest(input="Hello", voice="emb_voice", x_vector_only_mode=False) + params = speech_server._build_tts_params(req) assert params["x_vector_only_mode"] == [True] assert "voice_clone_prompt" in params @@ -1654,9 +1651,9 @@ class TestAsyncOmniSupportedTasks: assert "generate" in tasks -def test_api_server_create_speech_wraps_error_response_status(): - handler = MagicMock() - handler.create_speech = AsyncMock( +def test_api_server_create_speech_wraps_error_response_status(mocker: MockerFixture): + handler = mocker.MagicMock() + handler.create_speech = mocker.AsyncMock( return_value=ErrorResponse( error=ErrorInfo(message="bad request", type="BadRequestError", param=None, code=400), ) @@ -1851,9 +1848,9 @@ class TestFishSpeechServing: assert "<|speaker:0|>你好，[laughing]欢迎回来。<|speaker:1|>我也来了。" in encoded_texts assert all(allowed_special is None for _, _, allowed_special in tokenizer.calls) - def test_build_fish_clone_prompt_normalizes_text_fields(self, fish_speech_server): + def test_build_fish_clone_prompt_normalizes_text_fields(self, fish_speech_server, mocker: MockerFixture): fish_speech_server._fish_speech_tokenizer = _FakeFishTokenizer() - fish_speech_server._estimate_fish_prompt_len = MagicMock(return_value=123) + fish_speech_server._estimate_fish_prompt_len = mocker.MagicMock(return_value=123) request = OpenAICreateSpeechRequest( input="你好，欢迎回来。", @@ -1904,8 +1901,10 @@ class TestFishSpeechServing: with pytest.raises(ValueError, match="unsupported control token"): fish_speech_server._build_fish_speech_prompt(request) - def test_prepare_speech_generation_overrides_fish_default_max_tokens(self, fish_speech_server): - fish_speech_server._build_fish_speech_prompt_async = AsyncMock( + def test_prepare_speech_generation_overrides_fish_default_max_tokens( + self, fish_speech_server, mocker: MockerFixture + ): + fish_speech_server._build_fish_speech_prompt_async = mocker.AsyncMock( return_value={ "prompt_token_ids": [1, 2, 3], "additional_information": {}, @@ -1924,8 +1923,8 @@ class TestFishSpeechServing: assert sampling_params_list[0].max_tokens == 4096 assert fish_speech_server.engine_client.default_sampling_params_list[0].max_tokens == 2048 - def test_prepare_speech_generation_uses_stage_default_max_tokens(self, fish_speech_server): - fish_speech_server._build_fish_speech_prompt_async = AsyncMock( + def test_prepare_speech_generation_uses_stage_default_max_tokens(self, fish_speech_server, mocker: MockerFixture): + fish_speech_server._build_fish_speech_prompt_async = mocker.AsyncMock( return_value={ "prompt_token_ids": [1, 2, 3], "additional_information": {}, @@ -1956,9 +1955,9 @@ class TestFishSpeechServing: fish_speech_server.engine_client.generate.assert_not_called() - def test_create_speech_batch_allows_fish_text_only_items(self, fish_speech_server): - fish_speech_server._check_model = AsyncMock(return_value=None) - fish_speech_server._generate_audio_bytes = AsyncMock(return_value=("YWJj", "audio/wav")) + def test_create_speech_batch_allows_fish_text_only_items(self, fish_speech_server, mocker: MockerFixture): + fish_speech_server._check_model = mocker.AsyncMock(return_value=None) + fish_speech_server._generate_audio_bytes = mocker.AsyncMock(return_value=("YWJj", "audio/wav")) batch = BatchSpeechRequest(items=[SpeechBatchItem(input="hello fish")]) response = asyncio.run(fish_speech_server.create_speech_batch(batch)) @@ -2154,8 +2153,8 @@ class TestCosyVoice3Serving: assert error is not None assert "max_new_tokens" in error - def test_prepare_speech_generation_cosyvoice3(self, cosyvoice3_server): - cosyvoice3_server._build_cosyvoice3_prompt = AsyncMock( + def test_prepare_speech_generation_cosyvoice3(self, cosyvoice3_server, mocker: MockerFixture): + cosyvoice3_server._build_cosyvoice3_prompt = mocker.AsyncMock( return_value={ "prompt": "Hello", "multi_modal_data": {"audio": (np.zeros(24000), 24000)}, @@ -2236,9 +2235,9 @@ class TestTTSAsyncOffloading: yield server server.shutdown() - def test_prepare_speech_generation_awaits_voxtral_async(self, voxtral_server): + def test_prepare_speech_generation_awaits_voxtral_async(self, voxtral_server, mocker: MockerFixture): """Voxtral path in _prepare_speech_generation should call the async wrapper.""" - voxtral_server._build_voxtral_prompt_async = AsyncMock( + voxtral_server._build_voxtral_prompt_async = mocker.AsyncMock( return_value={ "prompt_token_ids": [1, 2, 3], "additional_information": {"voice": ["test"]}, @@ -2248,13 +2247,13 @@ class TestTTSAsyncOffloading: asyncio.run(voxtral_server._prepare_speech_generation(request)) voxtral_server._build_voxtral_prompt_async.assert_awaited_once() - def test_prepare_speech_generation_awaits_qwen3_tts_async(self, qwen3_tts_server): + def test_prepare_speech_generation_awaits_qwen3_tts_async(self, qwen3_tts_server, mocker: MockerFixture): """Qwen3 TTS path should call _estimate_prompt_len_async.""" - qwen3_tts_server._validate_tts_request = MagicMock(return_value=None) - qwen3_tts_server._build_tts_params = MagicMock( + qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None) + qwen3_tts_server._build_tts_params = mocker.MagicMock( return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]} ) - qwen3_tts_server._estimate_prompt_len_async = AsyncMock(return_value=512) + qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512) request = OpenAICreateSpeechRequest(input="hello") asyncio.run(qwen3_tts_server._prepare_speech_generation(request)) qwen3_tts_server._build_tts_params.assert_called_once() @@ -2281,8 +2280,8 @@ class TestTTSAsyncOffloading: server.shutdown() # Should not raise assert server._tts_executor is None - def test_diffusion_instance_shutdown_safe(self): + def test_diffusion_instance_shutdown_safe(self, mocker: MockerFixture): """Diffusion instances (created via for_diffusion) should have safe shutdown.""" - server = OmniOpenAIServingSpeech.for_diffusion(diffusion_engine=MagicMock(), model_name="test-model") + server = OmniOpenAIServingSpeech.for_diffusion(diffusion_engine=mocker.MagicMock(), model_name="test-model") assert server._tts_executor is None server.shutdown() # Should not raise diff --git a/tests/entrypoints/openai_api/test_serving_speech_stream.py b/tests/entrypoints/openai_api/test_serving_speech_stream.py index 1d26b5855..1b93ef58e 100644 --- a/tests/entrypoints/openai_api/test_serving_speech_stream.py +++ b/tests/entrypoints/openai_api/test_serving_speech_stream.py @@ -1,8 +1,8 @@ import asyncio -from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import FastAPI, WebSocket +from pytest_mock import MockerFixture from starlette.testclient import TestClient from starlette.websockets import WebSocketDisconnect @@ -13,19 +13,26 @@ from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpee pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -def _build_test_app(speech_service=None, *, idle_timeout=30.0, config_timeout=10.0): +def _build_test_app( + speech_service=None, + *, + idle_timeout=30.0, + config_timeout=10.0, + mocker: MockerFixture | None = None, +): if speech_service is None: - speech_service = MagicMock(spec=OmniOpenAIServingSpeech) - speech_service._generate_audio_bytes = AsyncMock(return_value=(b"RIFF" + b"\x00" * 32, "audio/wav")) - speech_service._prepare_speech_generation = AsyncMock(return_value=("req-1", object(), {})) + assert mocker is not None + speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) + speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"RIFF" + b"\x00" * 32, "audio/wav")) + speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {})) async def mock_generate_pcm_chunks(_generator, _request_id): for chunk in (b"\x01\x02", b"\x03\x04\x05"): yield chunk speech_service._generate_pcm_chunks = mock_generate_pcm_chunks - speech_service.engine_client = MagicMock() - speech_service.engine_client.abort = AsyncMock() + speech_service.engine_client = mocker.MagicMock() + speech_service.engine_client.abort = mocker.AsyncMock() handler = OmniStreamingSpeechHandler( speech_service=speech_service, @@ -42,8 +49,8 @@ def _build_test_app(speech_service=None, *, idle_timeout=30.0, config_timeout=10 class TestStreamingSpeechWebSocket: - def test_non_streaming_single_frame(self): - app, speech_service = _build_test_app() + def test_non_streaming_single_frame(self, mocker: MockerFixture): + app, speech_service = _build_test_app(mocker=mocker) with TestClient(app) as client: with client.websocket_connect("/v1/audio/speech/stream") as ws: @@ -68,13 +75,13 @@ class TestStreamingSpeechWebSocket: assert speech_service._generate_audio_bytes.await_count == 1 - def test_streaming_multiple_binary_frames(self): + def test_streaming_multiple_binary_frames(self, mocker: MockerFixture): captured_requests = [] - speech_service = MagicMock(spec=OmniOpenAIServingSpeech) - speech_service._generate_audio_bytes = AsyncMock(return_value=(b"", "audio/wav")) - speech_service.engine_client = MagicMock() - speech_service.engine_client.abort = AsyncMock() + speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) + speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav")) + speech_service.engine_client = mocker.MagicMock() + speech_service.engine_client.abort = mocker.AsyncMock() async def mock_prepare_speech_generation(request): captured_requests.append(request) @@ -123,8 +130,8 @@ class TestStreamingSpeechWebSocket: assert captured_requests[0].initial_codec_chunk_frames == 12 assert speech_service._generate_audio_bytes.await_count == 0 - def test_flush_on_input_done(self): - app, _ = _build_test_app() + def test_flush_on_input_done(self, mocker: MockerFixture): + app, _ = _build_test_app(mocker=mocker) with TestClient(app) as client: with client.websocket_connect("/v1/audio/speech/stream") as ws: @@ -142,8 +149,8 @@ class TestStreamingSpeechWebSocket: } assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} - def test_invalid_streaming_config(self): - app, _ = _build_test_app() + def test_invalid_streaming_config(self, mocker: MockerFixture): + app, _ = _build_test_app(mocker=mocker) with TestClient(app) as client: with client.websocket_connect("/v1/audio/speech/stream") as ws: @@ -159,8 +166,8 @@ class TestStreamingSpeechWebSocket: assert error["type"] == "error" assert "response_format='pcm'" in error["message"] - def test_empty_input_text_emits_no_audio(self): - app, speech_service = _build_test_app() + def test_empty_input_text_emits_no_audio(self, mocker: MockerFixture): + app, speech_service = _build_test_app(mocker=mocker) with TestClient(app) as client: with client.websocket_connect("/v1/audio/speech/stream") as ws: @@ -172,8 +179,8 @@ class TestStreamingSpeechWebSocket: assert speech_service._generate_audio_bytes.await_count == 0 - def test_multiple_sentences_increment_indices(self): - app, _ = _build_test_app() + def test_multiple_sentences_increment_indices(self, mocker: MockerFixture): + app, _ = _build_test_app(mocker=mocker) with TestClient(app) as client: with client.websocket_connect("/v1/audio/speech/stream") as ws: @@ -203,8 +210,8 @@ class TestStreamingSpeechWebSocket: ws.send_json({"type": "input.done"}) assert ws.receive_json() == {"type": "session.done", "total_sentences": 2} - def test_unknown_message_type_keeps_session_open(self): - app, _ = _build_test_app() + def test_unknown_message_type_keeps_session_open(self, mocker: MockerFixture): + app, _ = _build_test_app(mocker=mocker) with TestClient(app) as client: with client.websocket_connect("/v1/audio/speech/stream") as ws: @@ -227,21 +234,21 @@ class TestStreamingSpeechWebSocket: ws.send_json({"type": "input.done"}) assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} - def test_config_timeout_closes_session(self): - app, _ = _build_test_app(config_timeout=0.01) + def test_config_timeout_closes_session(self, mocker: MockerFixture): + app, _ = _build_test_app(config_timeout=0.01, mocker=mocker) with TestClient(app) as client: with client.websocket_connect("/v1/audio/speech/stream") as ws: error = ws.receive_json() assert error == {"type": "error", "message": "Timeout waiting for session.config"} - def test_generation_error_marks_audio_done(self): - speech_service = MagicMock(spec=OmniOpenAIServingSpeech) - speech_service._generate_audio_bytes = AsyncMock(side_effect=RuntimeError("boom")) - speech_service._prepare_speech_generation = AsyncMock(return_value=("req-err", object(), {})) - speech_service._generate_pcm_chunks = AsyncMock() - speech_service.engine_client = MagicMock() - speech_service.engine_client.abort = AsyncMock() + def test_generation_error_marks_audio_done(self, mocker: MockerFixture): + speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) + speech_service._generate_audio_bytes = mocker.AsyncMock(side_effect=RuntimeError("boom")) + speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-err", object(), {})) + speech_service._generate_pcm_chunks = mocker.AsyncMock() + speech_service.engine_client = mocker.MagicMock() + speech_service.engine_client.abort = mocker.AsyncMock() app, _ = _build_test_app(speech_service) with TestClient(app) as client: @@ -256,12 +263,12 @@ class TestStreamingSpeechWebSocket: ws.send_json({"type": "input.done"}) assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} - def test_streaming_generation_error_marks_audio_done(self): - speech_service = MagicMock(spec=OmniOpenAIServingSpeech) - speech_service._generate_audio_bytes = AsyncMock(return_value=(b"", "audio/wav")) - speech_service._prepare_speech_generation = AsyncMock(return_value=("req-stream-err", object(), {})) - speech_service.engine_client = MagicMock() - speech_service.engine_client.abort = AsyncMock() + def test_streaming_generation_error_marks_audio_done(self, mocker: MockerFixture): + speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) + speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav")) + speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-stream-err", object(), {})) + speech_service.engine_client = mocker.MagicMock() + speech_service.engine_client.abort = mocker.AsyncMock() async def mock_generate_pcm_chunks(_generator, _request_id): yield b"\x01\x02" @@ -298,8 +305,8 @@ class TestStreamingSpeechWebSocket: ws.send_json({"type": "input.done"}) assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} - def test_invalid_input_text_type_returns_validation_error(self): - app, speech_service = _build_test_app() + def test_invalid_input_text_type_returns_validation_error(self, mocker: MockerFixture): + app, speech_service = _build_test_app(mocker=mocker) with TestClient(app) as client: with client.websocket_connect("/v1/audio/speech/stream") as ws: @@ -316,9 +323,9 @@ class TestStreamingSpeechWebSocket: assert speech_service._generate_audio_bytes.await_count == 0 - def test_input_text_message_too_large(self, monkeypatch): + def test_input_text_message_too_large(self, monkeypatch, mocker: MockerFixture): monkeypatch.setattr(streaming_speech_module, "_MAX_INPUT_TEXT_MESSAGE_SIZE", 32) - app, speech_service = _build_test_app() + app, speech_service = _build_test_app(mocker=mocker) with TestClient(app) as client: with client.websocket_connect("/v1/audio/speech/stream") as ws: @@ -335,9 +342,9 @@ class TestStreamingSpeechWebSocket: assert speech_service._generate_audio_bytes.await_count == 0 - def test_session_config_message_too_large(self, monkeypatch): + def test_session_config_message_too_large(self, monkeypatch, mocker: MockerFixture): monkeypatch.setattr(streaming_speech_module, "_MAX_CONFIG_MESSAGE_SIZE", 64) - app, _ = _build_test_app() + app, _ = _build_test_app(mocker=mocker) with TestClient(app) as client: with client.websocket_connect("/v1/audio/speech/stream") as ws: @@ -348,12 +355,12 @@ class TestStreamingSpeechWebSocket: "message": "session.config message too large", } - def test_disconnect_aborts_streaming_request(self): - speech_service = MagicMock(spec=OmniOpenAIServingSpeech) - speech_service._generate_audio_bytes = AsyncMock(return_value=(b"", "audio/wav")) - speech_service._prepare_speech_generation = AsyncMock(return_value=("req-abort", object(), {})) - speech_service.engine_client = MagicMock() - speech_service.engine_client.abort = AsyncMock() + def test_disconnect_aborts_streaming_request(self, mocker: MockerFixture): + speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) + speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav")) + speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-abort", object(), {})) + speech_service.engine_client = mocker.MagicMock() + speech_service.engine_client.abort = mocker.AsyncMock() async def mock_generate_pcm_chunks(_generator, _request_id): yield b"\x01\x02" @@ -361,11 +368,11 @@ class TestStreamingSpeechWebSocket: speech_service._generate_pcm_chunks = mock_generate_pcm_chunks handler = OmniStreamingSpeechHandler(speech_service=speech_service) - websocket = MagicMock() - websocket.send_json = AsyncMock(side_effect=[None, WebSocketDisconnect()]) - websocket.send_bytes = AsyncMock(side_effect=WebSocketDisconnect()) + websocket = mocker.MagicMock() + websocket.send_json = mocker.AsyncMock(side_effect=[None, WebSocketDisconnect()]) + websocket.send_bytes = mocker.AsyncMock(side_effect=WebSocketDisconnect()) - config = MagicMock() + config = mocker.MagicMock() config.model = None config.voice = "Vivian" config.task_type = None diff --git a/tests/entrypoints/test_omni_base_profiler.py b/tests/entrypoints/test_omni_base_profiler.py index 0c1ddc6a5..ca10eed91 100644 --- a/tests/entrypoints/test_omni_base_profiler.py +++ b/tests/entrypoints/test_omni_base_profiler.py @@ -1,8 +1,7 @@ """Unit tests for OmniBase and AsyncOmni profiler methods.""" -from unittest.mock import MagicMock, patch - import pytest +from pytest_mock import MockerFixture pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -11,12 +10,12 @@ class TestOmniBaseProfiler: """Test suite for OmniBase profiler methods (start_profile, stop_profile).""" @pytest.fixture - def mock_engine(self): + def mock_engine(self, mocker: MockerFixture): """Create a mock AsyncOmniEngine for testing.""" - engine = MagicMock() + engine = mocker.MagicMock() engine.num_stages = 3 engine.is_alive.return_value = True - engine.default_sampling_params_list = [MagicMock() for _ in range(3)] + engine.default_sampling_params_list = [mocker.MagicMock() for _ in range(3)] engine.get_stage_metadata.side_effect = lambda i: { "final_output_type": "text" if i == 0 else "audio", "final_output": True, @@ -25,17 +24,15 @@ class TestOmniBaseProfiler: return engine @pytest.fixture - def omni_base_instance(self, mock_engine): + def omni_base_instance(self, mock_engine, mocker: MockerFixture): """Create an OmniBase instance with mocked dependencies.""" - with ( - patch("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", return_value=mock_engine), - patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", side_effect=lambda x: x), - patch("vllm_omni.entrypoints.omni_base.weakref.finalize"), - ): - from vllm_omni.entrypoints.omni_base import OmniBase + mocker.patch("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", return_value=mock_engine) + mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", side_effect=lambda x: x) + mocker.patch("vllm_omni.entrypoints.omni_base.weakref.finalize") + from vllm_omni.entrypoints.omni_base import OmniBase - instance = OmniBase(model="test-model") - return instance + instance = OmniBase(model="test-model") + return instance def test_start_profile_calls_collective_rpc(self, omni_base_instance, mock_engine): """Test that start_profile calls collective_rpc with correct arguments.""" diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 916db3cc2..afa7fa82e 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -3,9 +3,9 @@ from __future__ import annotations import argparse -from unittest.mock import Mock, patch import pytest +from pytest_mock import MockerFixture from vllm_omni.entrypoints.cli.serve import run_headless @@ -26,45 +26,43 @@ def _make_headless_args() -> argparse.Namespace: ) -def test_run_headless_registers_stage_once_and_launches_all_local_engines() -> None: +def test_run_headless_registers_stage_once_and_launches_all_local_engines(mocker: MockerFixture) -> None: args = _make_headless_args() - stage_cfg = Mock(stage_id=3) + stage_cfg = mocker.Mock(stage_id=3) stage_cfgs = [stage_cfg] - parallel_config = Mock( + parallel_config = mocker.Mock( data_parallel_size_local=2, data_parallel_rank=4, data_parallel_rank_local=1, node_rank_within_dp=0, ) - vllm_config = Mock(parallel_config=parallel_config) - executor_class = Mock() - engine_manager = Mock() + vllm_config = mocker.Mock(parallel_config=parallel_config) + executor_class = mocker.Mock() + engine_manager = mocker.Mock() - with ( - patch( - "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", stage_cfgs), - ), - patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment"), - patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=Mock()), - patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}), - patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}), - patch( - "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage", - return_value=(None, None, None), - ), - patch( - "vllm_omni.engine.stage_init_utils.build_vllm_config", - return_value=(vllm_config, executor_class), - ) as mock_build_vllm_config, - patch( - "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", - return_value="tcp://127.0.0.1:26001", - ) as mock_register, - patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager) as mock_manager_cls, - patch("signal.signal"), - ): - run_headless(args) + mocker.patch( + "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", + return_value=("/fake/stages.yaml", stage_cfgs), + ) + mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") + mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=mocker.Mock()) + mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}) + mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}) + mocker.patch( + "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ) + mock_build_vllm_config = mocker.patch( + "vllm_omni.engine.stage_init_utils.build_vllm_config", + return_value=(vllm_config, executor_class), + ) + mock_register = mocker.patch( + "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", + return_value="tcp://127.0.0.1:26001", + ) + mock_manager_cls = mocker.patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager) + mocker.patch("signal.signal") + run_headless(args) mock_build_vllm_config.assert_called_once_with( stage_cfg, @@ -92,89 +90,85 @@ def test_run_headless_registers_stage_once_and_launches_all_local_engines() -> N engine_manager.shutdown.assert_called_once_with() -def test_run_headless_honors_explicit_log_stats_flag() -> None: +def test_run_headless_honors_explicit_log_stats_flag(mocker: MockerFixture) -> None: args = _make_headless_args() args.log_stats = True - stage_cfg = Mock(stage_id=3) + stage_cfg = mocker.Mock(stage_id=3) stage_cfgs = [stage_cfg] - parallel_config = Mock( + parallel_config = mocker.Mock( data_parallel_size_local=2, data_parallel_rank=4, data_parallel_rank_local=1, node_rank_within_dp=0, ) - vllm_config = Mock(parallel_config=parallel_config) - executor_class = Mock() - engine_manager = Mock() + vllm_config = mocker.Mock(parallel_config=parallel_config) + executor_class = mocker.Mock() + engine_manager = mocker.Mock() - with ( - patch( - "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", stage_cfgs), - ), - patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment"), - patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=Mock()), - patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}), - patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}), - patch( - "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage", - return_value=(None, None, None), - ), - patch( - "vllm_omni.engine.stage_init_utils.build_vllm_config", - return_value=(vllm_config, executor_class), - ), - patch( - "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", - return_value="tcp://127.0.0.1:26001", - ), - patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager) as mock_manager_cls, - patch("signal.signal"), - ): - run_headless(args) + mocker.patch( + "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", + return_value=("/fake/stages.yaml", stage_cfgs), + ) + mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") + mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=mocker.Mock()) + mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}) + mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}) + mocker.patch( + "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ) + mocker.patch( + "vllm_omni.engine.stage_init_utils.build_vllm_config", + return_value=(vllm_config, executor_class), + ) + mocker.patch( + "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", + return_value="tcp://127.0.0.1:26001", + ) + mock_manager_cls = mocker.patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager) + mocker.patch("signal.signal") + run_headless(args) manager_kwargs = mock_manager_cls.call_args.kwargs assert manager_kwargs["log_stats"] is True -def test_run_headless_launches_diffusion_stage_via_omni_master() -> None: +def test_run_headless_launches_diffusion_stage_via_omni_master(mocker: MockerFixture) -> None: args = _make_headless_args() - stage_cfg = Mock(stage_id=3, stage_type="diffusion") - stage_cfg.engine_args = Mock() + stage_cfg = mocker.Mock(stage_id=3, stage_type="diffusion") + stage_cfg.engine_args = mocker.Mock() stage_cfg.engine_input_source = [] stage_cfgs = [stage_cfg] - metadata = Mock(stage_id=3) - od_config = Mock() - proc = Mock() + metadata = mocker.Mock(stage_id=3) + od_config = mocker.Mock() + proc = mocker.Mock() proc.exitcode = 0 proc.is_alive.return_value = False - with ( - patch( - "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", stage_cfgs), - ), - patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment"), - patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=Mock()), - patch( - "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage", - return_value=(None, None, None), - ), - patch("vllm_omni.engine.stage_init_utils.extract_stage_metadata", return_value=metadata), - patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info") as mock_inject_stage_info, - patch("vllm_omni.engine.stage_init_utils.build_diffusion_config", return_value=od_config), - patch( - "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", - return_value=("tcp://127.0.0.1:26001", "tcp://127.0.0.1:26002", "tcp://127.0.0.1:26003"), - ) as mock_register, - patch( - "vllm_omni.diffusion.stage_diffusion_proc.spawn_diffusion_proc", - return_value=(proc, "tcp://127.0.0.1:26001", "tcp://127.0.0.1:26002", "tcp://127.0.0.1:26003"), - ) as mock_spawn, - patch("vllm_omni.diffusion.stage_diffusion_proc.complete_diffusion_handshake") as mock_handshake, - patch("signal.signal"), - ): - run_headless(args) + mocker.patch( + "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", + return_value=("/fake/stages.yaml", stage_cfgs), + ) + mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") + mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=mocker.Mock()) + mocker.patch( + "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ) + mocker.patch("vllm_omni.engine.stage_init_utils.extract_stage_metadata", return_value=metadata) + mock_inject_stage_info = mocker.patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info") + mocker.patch("vllm_omni.engine.stage_init_utils.build_diffusion_config", return_value=od_config) + mock_register = mocker.patch( + "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", + return_value=("tcp://127.0.0.1:26001", "tcp://127.0.0.1:26002", "tcp://127.0.0.1:26003"), + ) + mock_spawn = mocker.patch( + "vllm_omni.diffusion.stage_diffusion_proc.spawn_diffusion_proc", + return_value=(proc, "tcp://127.0.0.1:26001", "tcp://127.0.0.1:26002", "tcp://127.0.0.1:26003"), + ) + mock_handshake = mocker.patch("vllm_omni.diffusion.stage_diffusion_proc.complete_diffusion_handshake") + mocker.patch("signal.signal") + run_headless(args) mock_inject_stage_info.assert_called_once_with(stage_cfg, 3) mock_register.assert_called_once_with( diff --git a/tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py b/tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py index 85c0e8b56..8858d1f8f 100644 --- a/tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py +++ b/tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py @@ -2,10 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace -from unittest.mock import Mock import pytest import torch +from pytest_mock import MockerFixture from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import TALKER_CODEC_PAD_TOKEN_ID from vllm_omni.model_executor.models.mimo_audio.mimo_audio_code2wav import ( @@ -51,7 +51,7 @@ def _make_invalid_flat_immediate_eostm(eostm_id: int = 666) -> torch.Tensor: return g.reshape(-1) -def _minimal_model(): +def _minimal_model(mocker: MockerFixture): """Avoid __init__ (HF tokenizer paths); only fields used by _batch_decode_waveforms.""" model = object.__new__(MiMoAudioToken2WavForConditionalGenerationVLLM) model.device = torch.device("cpu") @@ -59,7 +59,7 @@ def _minimal_model(): model.streamer_config = AudioStreamerConfig(group_size=_GROUP, audio_channels=_AC) model.codes = _codes_ns() - decode_vq = Mock( + decode_vq = mocker.Mock( side_effect=lambda audio_codes: torch.ones( audio_codes.shape[1], 7, @@ -67,7 +67,7 @@ def _minimal_model(): device=audio_codes.device, ) ) - decoder = Mock() + decoder = mocker.Mock() audio_tok = SimpleNamespace( encoder=SimpleNamespace(decode_vq=decode_vq), @@ -78,9 +78,9 @@ def _minimal_model(): return model, audio_tok -def test_batch_decode_waveforms_empty_input_list(): +def test_batch_decode_waveforms_empty_input_list(mocker: MockerFixture): """Empty input list returns a single zero-length float32 tensor on model device.""" - model, _ = _minimal_model() + model, _ = _minimal_model(mocker) out = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(model, []) assert len(out) == 1 assert out[0].dtype == torch.float32 @@ -88,9 +88,9 @@ def test_batch_decode_waveforms_empty_input_list(): assert out[0].device == model.device -def test_batch_decode_waveforms_single_vs_multiple_decoder_shapes(): +def test_batch_decode_waveforms_single_vs_multiple_decoder_shapes(mocker: MockerFixture): """Single and multi-request batches produce correctly shaped packed hidden states and trimmed waveforms.""" - model, audio_tok = _minimal_model() + model, audio_tok = _minimal_model(mocker) decoder = audio_tok.decoder # Single valid request: decoder output rank-3 for double squeeze path @@ -118,9 +118,9 @@ def test_batch_decode_waveforms_single_vs_multiple_decoder_shapes(): assert out2[1].shape == (8 * _FTP,) -def test_batch_decode_waveforms_mixed_valid_invalid_requests(): +def test_batch_decode_waveforms_mixed_valid_invalid_requests(mocker: MockerFixture): """Mixed valid and invalid requests: invalid slots get empty tensors, valid slots get decoded waveforms.""" - model, audio_tok = _minimal_model() + model, audio_tok = _minimal_model(mocker) valid_a = _make_valid_flat_codes(1) valid_b = _make_valid_flat_codes(1) dummy = _make_dummy_code_tensor() @@ -151,9 +151,9 @@ def test_batch_decode_waveforms_mixed_valid_invalid_requests(): assert input_lengths.tolist() == [4, 4] -def test_batch_decode_waveforms_all_invalid_returns_per_request_empty(): +def test_batch_decode_waveforms_all_invalid_returns_per_request_empty(mocker: MockerFixture): """All-invalid batch skips decoder entirely and returns empty tensors for every slot.""" - model, audio_tok = _minimal_model() + model, audio_tok = _minimal_model(mocker) out = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms( model, [None, _make_dummy_code_tensor(), torch.tensor([], dtype=torch.long)], @@ -163,9 +163,9 @@ def test_batch_decode_waveforms_all_invalid_returns_per_request_empty(): audio_tok.decoder.assert_not_called() -def test_batch_decode_waveforms_output_shape_trim_when_decoder_returns_extra_samples(): +def test_batch_decode_waveforms_output_shape_trim_when_decoder_returns_extra_samples(mocker: MockerFixture): """Decoder output longer than valid_len is trimmed to the exact expected waveform length.""" - model, audio_tok = _minimal_model() + model, audio_tok = _minimal_model(mocker) flat = _make_valid_flat_codes(1) # Longer than valid_len so branch wav = wav[:valid_len] runs audio_tok.decoder.return_value = torch.ones(1, 1, 10_000, dtype=torch.float32) @@ -175,9 +175,9 @@ def test_batch_decode_waveforms_output_shape_trim_when_decoder_returns_extra_sam assert out[0].dtype == torch.float32 -def test_batch_decode_waveforms_multi_request_trims_each_row_when_decoder_returns_extra(): +def test_batch_decode_waveforms_multi_request_trims_each_row_when_decoder_returns_extra(mocker: MockerFixture): """Else-branch split: per-request wav[:valid_len] when decoder pads each batch row.""" - model, audio_tok = _minimal_model() + model, audio_tok = _minimal_model(mocker) a = _make_valid_flat_codes(1) b = _make_valid_flat_codes(2) audio_tok.decoder.return_value = torch.ones(2, 1, 10_000, dtype=torch.float32) @@ -189,9 +189,9 @@ def test_batch_decode_waveforms_multi_request_trims_each_row_when_decoder_return assert out[1].dtype == torch.float32 -def test_batch_decode_waveforms_valid_only_at_edges_maps_to_correct_indices(): +def test_batch_decode_waveforms_valid_only_at_edges_maps_to_correct_indices(mocker: MockerFixture): """Tensor packing order must match valid_indices when invalid requests are in the middle.""" - model, audio_tok = _minimal_model() + model, audio_tok = _minimal_model(mocker) first = _make_valid_flat_codes(1) last = _make_valid_flat_codes(2) inputs = [ @@ -212,9 +212,9 @@ def test_batch_decode_waveforms_valid_only_at_edges_maps_to_correct_indices(): assert input_lengths.tolist() == [4, 8] -def test_batch_decode_waveforms_output_shapes_1d_float32_for_all_slots(): +def test_batch_decode_waveforms_output_shapes_1d_float32_for_all_slots(mocker: MockerFixture): """Every slot is a 1-D float32 vector (empty or waveform), matching downstream expectations.""" - model, audio_tok = _minimal_model() + model, audio_tok = _minimal_model(mocker) inputs = [_make_valid_flat_codes(1), None, _make_valid_flat_codes(1)] audio_tok.decoder.return_value = torch.ones(2, 1, 5000, dtype=torch.float32) out = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(model, inputs) diff --git a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py index 8e04b0496..587e7f7f8 100644 --- a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py +++ b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py @@ -10,10 +10,9 @@ Regression test for: https://github.com/vllm-project/vllm/issues/34506 - Interleaved (use_audio_in_video) should also work correctly. """ -from unittest.mock import Mock - import pytest import torch +from pytest_mock import MockerFixture from vllm.model_executor.models.qwen2_5_omni_thinker import ( check_interleaved_audio_video, merge_interleaved_embeddings, @@ -107,7 +106,7 @@ class TestCheckInterleavedAudioVideo: # --------------------------------------------------------------------------- -def make_mock_model(hidden: int = 8): +def make_mock_model(mocker: MockerFixture, hidden: int = 8): """ Return a minimal mock of Qwen2_5OmniThinkerForConditionalGeneration that has enough structure to run embed_input_ids. @@ -116,10 +115,10 @@ def make_mock_model(hidden: int = 8): Qwen2_5OmniThinkerForConditionalGeneration, ) - model = Mock(spec=Qwen2_5OmniThinkerForConditionalGeneration) + model = mocker.Mock(spec=Qwen2_5OmniThinkerForConditionalGeneration) # Config with token IDs - cfg = Mock() + cfg = mocker.Mock() cfg.video_token_index = VIDEO_TOKEN_ID cfg.audio_token_index = AUDIO_TOKEN_ID model.config = cfg @@ -130,9 +129,9 @@ def make_mock_model(hidden: int = 8): # view with shared memory, which masked_scatter_ cannot handle). return ids.float().unsqueeze(-1).expand(-1, hidden).clone() - lang_model = Mock() + lang_model = mocker.Mock() lang_model.embed_input_ids = fake_lm_embed - model.get_language_model = Mock(return_value=lang_model) + model.get_language_model = mocker.Mock(return_value=lang_model) from vllm.model_executor.models.interfaces import SupportsMultiModal @@ -169,7 +168,7 @@ def build_mm_embeds(audio_n, image_n, video_n, hidden, audio_val=10.0, image_val class TestEmbedInputIds: - def _run(self, audio_n, image_n, video_n, hidden=8): + def _run(self, mocker: MockerFixture, audio_n, image_n, video_n, hidden=8): """ Run embed_input_ids for a non-interleaved mixed-modality sequence. Returns (result_embeds, input_ids, is_multimodal). @@ -177,33 +176,33 @@ class TestEmbedInputIds: input_ids, is_multimodal = make_token_seq(audio_n, image_n, video_n) mm_embeds = build_mm_embeds(audio_n, image_n, video_n, hidden) - model, _ = make_mock_model(hidden) + model, _ = make_mock_model(mocker, hidden) result = model.embed_input_ids(input_ids, mm_embeds, is_multimodal=is_multimodal) return result, input_ids, is_multimodal - def test_audio_only(self): + def test_audio_only(self, mocker: MockerFixture): """Audio-only: audio positions get audio embeddings.""" audio_n, hidden = 5, 8 audio_val = 10.0 - result, input_ids, is_multimodal = self._run(audio_n, 0, 0, hidden) + result, input_ids, is_multimodal = self._run(mocker, audio_n, 0, 0, hidden) audio_pos = (input_ids == AUDIO_TOKEN_ID).nonzero(as_tuple=True)[0] assert result[audio_pos].allclose(torch.full((audio_n, hidden), audio_val)), ( "Audio positions should get audio embeddings" ) - def test_video_only(self): + def test_video_only(self, mocker: MockerFixture): """Video-only: video positions get video embeddings.""" video_n, hidden = 6, 8 video_val = 30.0 - result, input_ids, is_multimodal = self._run(0, 0, video_n, hidden) + result, input_ids, is_multimodal = self._run(mocker, 0, 0, video_n, hidden) video_pos = (input_ids == VIDEO_TOKEN_ID).nonzero(as_tuple=True)[0] assert result[video_pos].allclose(torch.full((video_n, hidden), video_val)), ( "Video positions should get video embeddings" ) - def test_mixed_modalities_audio_goes_to_audio_pos(self): + def test_mixed_modalities_audio_goes_to_audio_pos(self, mocker: MockerFixture): """ Regression test for GitHub issue #34506: With audio + image + video (non-interleaved), audio positions must @@ -212,7 +211,7 @@ class TestEmbedInputIds: audio_n, image_n, video_n, hidden = 5, 4, 6, 8 audio_val, image_val, video_val = 10.0, 20.0, 30.0 - result, input_ids, is_multimodal = self._run(audio_n, image_n, video_n, hidden) + result, input_ids, is_multimodal = self._run(mocker, audio_n, image_n, video_n, hidden) audio_pos = (input_ids == AUDIO_TOKEN_ID).nonzero(as_tuple=True)[0] image_pos = (input_ids == IMAGE_TOKEN_ID).nonzero(as_tuple=True)[0] @@ -233,10 +232,10 @@ class TestEmbedInputIds: f"Video emb wrong: expected {video_val}, got mean={mean_v:.1f}" ) - def test_text_positions_unchanged(self): + def test_text_positions_unchanged(self, mocker: MockerFixture): """Text positions should keep their text embeddings.""" audio_n, image_n, video_n, hidden = 3, 2, 4, 8 - result, input_ids, is_multimodal = self._run(audio_n, image_n, video_n, hidden) + result, input_ids, is_multimodal = self._run(mocker, audio_n, image_n, video_n, hidden) text_pos = (~is_multimodal).nonzero(as_tuple=True)[0] # Text tokens have value TEXT_TOKEN_ID=0, so embed -> 0.0 @@ -244,7 +243,7 @@ class TestEmbedInputIds: "Text positions should keep text embeddings" ) - def test_interleaved_use_audio_in_video(self): + def test_interleaved_use_audio_in_video(self, mocker: MockerFixture): """ Interleaved (use_audio_in_video): video chunks interleaved with audio. Video embeddings must go to video positions, audio to audio positions. @@ -263,7 +262,7 @@ class TestEmbedInputIds: torch.full((audio_n, hidden), audio_val), ] - model, _ = make_mock_model(hidden) + model, _ = make_mock_model(mocker, hidden) result = model.embed_input_ids(input_ids, mm_embeds, is_multimodal=is_multimodal) video_pos = (input_ids == VIDEO_TOKEN_ID).nonzero(as_tuple=True)[0] diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py index e2970dcb2..b0ce10a8d 100644 --- a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py +++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py @@ -15,9 +15,10 @@ import importlib.util import os import sys import types -from unittest.mock import MagicMock, patch +import pytest import torch +from pytest_mock import MockerFixture # Direct file import to avoid vllm_omni.__init__ patch dependencies. _BASE = os.path.join( @@ -41,28 +42,31 @@ def _load_module(name: str, filename: str): return mod -def _build_mock_modules() -> dict[str, object]: +def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]: """Build the dict of modules to inject into sys.modules.""" - platforms_mock = MagicMock() + platforms_mock = mocker.MagicMock() platforms_mock.current_omni_platform.supports_torch_inductor.return_value = False - logger_mock = MagicMock() - logger_mock.init_logger = lambda name: MagicMock() + logger_mock = mocker.MagicMock() + logger_mock.init_logger = lambda name: mocker.MagicMock() - vllm_config_mod = MagicMock() - vllm_config_mod.set_current_vllm_config = lambda cfg: MagicMock(__enter__=MagicMock(), __exit__=MagicMock()) + vllm_config_mod = mocker.MagicMock() + vllm_config_mod.set_current_vllm_config = lambda cfg: mocker.MagicMock( + __enter__=mocker.MagicMock(), + __exit__=mocker.MagicMock(), + ) - weight_utils_mock = MagicMock() + weight_utils_mock = mocker.MagicMock() weight_utils_mock.default_weight_loader = lambda p, w: None pkg = types.ModuleType("vllm_omni.model_executor.models.qwen3_tts") pkg.__path__ = [os.path.abspath(_BASE)] return { - "vllm_omni": MagicMock(), + "vllm_omni": mocker.MagicMock(), "vllm_omni.platforms": platforms_mock, "vllm.logger": logger_mock, - "vllm.config": MagicMock(), + "vllm.config": mocker.MagicMock(), "vllm.config.vllm": vllm_config_mod, "vllm.model_executor.model_loader.weight_utils": weight_utils_mock, "vllm_omni.model_executor": types.ModuleType("vllm_omni.model_executor"), @@ -71,38 +75,47 @@ def _build_mock_modules() -> dict[str, object]: } -def _load_target_classes(): +def _load_target_classes(mocker: MockerFixture): """Load config and code predictor modules with mocked dependencies. - Uses patch.dict to ensure sys.modules is always restored, even on failure. + Uses mocker.patch.dict to ensure sys.modules is always restored, even on failure. """ - mocks = _build_mock_modules() - with patch.dict(sys.modules, mocks): - config_mod = _load_module( - "vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts", - "configuration_qwen3_tts.py", - ) - sys.modules["vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts"] = config_mod + mocks = _build_mock_modules(mocker) + mocker.patch.dict(sys.modules, mocks) + config_mod = _load_module( + "vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts", + "configuration_qwen3_tts.py", + ) + sys.modules["vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts"] = config_mod - cp_mod = _load_module( - "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm", - "qwen3_tts_code_predictor_vllm.py", - ) + cp_mod = _load_module( + "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm", + "qwen3_tts_code_predictor_vllm.py", + ) return config_mod, cp_mod -_config_mod, _cp_mod = _load_target_classes() - -Qwen3TTSTalkerCodePredictorConfig = _config_mod.Qwen3TTSTalkerCodePredictorConfig -Qwen3TTSTalkerConfig = _config_mod.Qwen3TTSTalkerConfig -CodePredictorWrapper = _cp_mod.Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM -CodePredictorModel = _cp_mod.Qwen3TTSTalkerCodePredictorModelVLLM +@pytest.fixture +def loaded_target_classes(mocker: MockerFixture): + config_mod, cp_mod = _load_target_classes(mocker) + return ( + config_mod.Qwen3TTSTalkerCodePredictorConfig, + config_mod.Qwen3TTSTalkerConfig, + cp_mod.Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM, + cp_mod.Qwen3TTSTalkerCodePredictorModelVLLM, + ) -def _make_tiny_config() -> tuple: +def _make_tiny_config(loaded_target_classes) -> tuple: """Create minimal configs for a tiny code predictor model.""" - cp_config = Qwen3TTSTalkerCodePredictorConfig( + ( + qwen3_tts_talker_code_predictor_config, + qwen3_tts_talker_config, + _, + _, + ) = loaded_target_classes + cp_config = qwen3_tts_talker_code_predictor_config( vocab_size=64, hidden_size=32, intermediate_size=64, @@ -113,16 +126,16 @@ def _make_tiny_config() -> tuple: num_code_groups=4, rms_norm_eps=1e-6, ) - talker_config = Qwen3TTSTalkerConfig( + talker_config = qwen3_tts_talker_config( hidden_size=32, num_code_groups=4, ) return cp_config, talker_config -def _make_vllm_config(max_num_seqs: int = 4) -> MagicMock: +def _make_vllm_config(mocker: MockerFixture, max_num_seqs: int = 4): """Create a mock VllmConfig with scheduler_config.""" - vllm_config = MagicMock() + vllm_config = mocker.MagicMock() vllm_config.scheduler_config.max_num_seqs = max_num_seqs return vllm_config @@ -130,12 +143,13 @@ def _make_vllm_config(max_num_seqs: int = 4) -> MagicMock: class TestCodePredictorDtypeAlignment: """Test that code predictor buffers match model parameter dtype.""" - def test_ensure_buffers_uses_given_dtype(self) -> None: + def test_ensure_buffers_uses_given_dtype(self, mocker: MockerFixture, loaded_target_classes) -> None: """_ensure_buffers should create proj_buf with the given dtype.""" - cp_config, talker_config = _make_tiny_config() - vllm_config = _make_vllm_config() + _, _, code_predictor_wrapper, _ = loaded_target_classes + cp_config, talker_config = _make_tiny_config(loaded_target_classes) + vllm_config = _make_vllm_config(mocker) - predictor = CodePredictorWrapper( + predictor = code_predictor_wrapper( vllm_config=vllm_config, config=cp_config, talker_config=talker_config, @@ -150,12 +164,13 @@ class TestCodePredictorDtypeAlignment: predictor._ensure_buffers(torch.device("cpu"), torch.float32) assert predictor._proj_buf.dtype == torch.float32 - def test_warmup_aligns_buffer_to_model_params(self) -> None: + def test_warmup_aligns_buffer_to_model_params(self, mocker: MockerFixture, loaded_target_classes) -> None: """_warmup_buckets should align proj_buf dtype to model parameters.""" - cp_config, talker_config = _make_tiny_config() - vllm_config = _make_vllm_config(max_num_seqs=2) + _, _, code_predictor_wrapper, _ = loaded_target_classes + cp_config, talker_config = _make_tiny_config(loaded_target_classes) + vllm_config = _make_vllm_config(mocker, max_num_seqs=2) - predictor = CodePredictorWrapper( + predictor = code_predictor_wrapper( vllm_config=vllm_config, config=cp_config, talker_config=talker_config, @@ -177,12 +192,13 @@ class TestCodePredictorDtypeAlignment: assert predictor._proj_buf.dtype == torch.float16 - def test_setup_compile_caches_model_dtype(self) -> None: + def test_setup_compile_caches_model_dtype(self, mocker: MockerFixture, loaded_target_classes) -> None: """_setup_compile should cache model parameter dtype.""" - cp_config, talker_config = _make_tiny_config() - vllm_config = _make_vllm_config(max_num_seqs=2) + _, _, code_predictor_wrapper, _ = loaded_target_classes + cp_config, talker_config = _make_tiny_config(loaded_target_classes) + vllm_config = _make_vllm_config(mocker, max_num_seqs=2) - predictor = CodePredictorWrapper( + predictor = code_predictor_wrapper( vllm_config=vllm_config, config=cp_config, talker_config=talker_config, @@ -193,12 +209,13 @@ class TestCodePredictorDtypeAlignment: predictor._setup_compile() assert predictor._model_dtype == torch.float16 - def test_forward_with_mismatched_input_dtype(self) -> None: + def test_forward_with_mismatched_input_dtype(self, mocker: MockerFixture, loaded_target_classes) -> None: """forward() should not crash when inputs are float32 but model is float16.""" - cp_config, talker_config = _make_tiny_config() - vllm_config = _make_vllm_config(max_num_seqs=2) + _, _, code_predictor_wrapper, _ = loaded_target_classes + cp_config, talker_config = _make_tiny_config(loaded_target_classes) + vllm_config = _make_vllm_config(mocker, max_num_seqs=2) - predictor = CodePredictorWrapper( + predictor = code_predictor_wrapper( vllm_config=vllm_config, config=cp_config, talker_config=talker_config, @@ -231,10 +248,11 @@ class TestCodePredictorDtypeAlignment: class TestCodePredictorModelDtype: """Test the inner model forward with different dtypes.""" - def test_model_forward_float16(self) -> None: + def test_model_forward_float16(self, loaded_target_classes) -> None: """Inner model forward should work in float16.""" - cp_config, _ = _make_tiny_config() - model = CodePredictorModel(cp_config, talker_hidden_size=32).to(torch.float16) + _, _, _, code_predictor_model = loaded_target_classes + cp_config, _ = _make_tiny_config(loaded_target_classes) + model = code_predictor_model(cp_config, talker_hidden_size=32).to(torch.float16) bsz, seq_len = 1, 4 inputs = torch.randn(bsz, seq_len, 32, dtype=torch.float16) @@ -244,10 +262,11 @@ class TestCodePredictorModelDtype: assert output.dtype == torch.float16 assert output.shape == (bsz, seq_len, 32) - def test_model_forward_float32(self) -> None: + def test_model_forward_float32(self, loaded_target_classes) -> None: """Inner model forward should work in float32.""" - cp_config, _ = _make_tiny_config() - model = CodePredictorModel(cp_config, talker_hidden_size=32).to(torch.float32) + _, _, _, code_predictor_model = loaded_target_classes + cp_config, _ = _make_tiny_config(loaded_target_classes) + model = code_predictor_model(cp_config, talker_hidden_size=32).to(torch.float32) bsz, seq_len = 1, 4 inputs = torch.randn(bsz, seq_len, 32, dtype=torch.float32) diff --git a/tests/model_executor/models/test_fish_speech_voice_cache.py b/tests/model_executor/models/test_fish_speech_voice_cache.py index 8fe7a4a4d..fef4b551a 100644 --- a/tests/model_executor/models/test_fish_speech_voice_cache.py +++ b/tests/model_executor/models/test_fish_speech_voice_cache.py @@ -10,11 +10,11 @@ Covers: import os import tempfile -from unittest.mock import MagicMock, patch import numpy as np import pytest import torch +from pytest_mock import MockerFixture pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -61,18 +61,18 @@ class TestFishSpeechVoiceCacheIntegration: """Test the cache-hit / cache-miss / no-cache paths in the model.""" @pytest.fixture - def mock_model(self): + def mock_model(self, mocker: MockerFixture): """Create a mock FishSpeechSlowARForConditionalGeneration with cache.""" from vllm_omni.utils.voice_cache import VoiceEmbeddingCache - model = MagicMock() + model = mocker.MagicMock() model._voice_cache = VoiceEmbeddingCache(max_entries=4) model._semantic_begin_id = 151678 model._num_codebooks = 10 model._codebook_size = 4096 model.model_path = "/fake/model" - model.codebook_embeddings = MagicMock() - model.codebook_embeddings.weight = MagicMock() + model.codebook_embeddings = mocker.MagicMock() + model.codebook_embeddings.weight = mocker.MagicMock() model.codebook_embeddings.weight.device = torch.device("cpu") return model @@ -166,9 +166,9 @@ class TestFishSpeechVoiceCacheIntegration: class TestFishSpeechValidatorUploadedVoice: """Test _validate_fish_tts_request uploaded voice resolution.""" - def test_uploaded_voice_resolves_ref_audio(self): + def test_uploaded_voice_resolves_ref_audio(self, mocker: MockerFixture): """When voice matches an uploaded speaker, ref_audio should be auto-set.""" - request = MagicMock() + request = mocker.MagicMock() request.input = "Hello" request.voice = "alice" request.ref_audio = None @@ -185,17 +185,17 @@ class TestFishSpeechValidatorUploadedVoice: } # Simulate: voice in uploaded_speakers, file exists, get_audio returns data URL. - with patch("pathlib.Path.exists", return_value=True): - voice_lower = request.voice.lower() - assert voice_lower in uploaded_speakers + mocker.patch("pathlib.Path.exists", return_value=True) + voice_lower = request.voice.lower() + assert voice_lower in uploaded_speakers - speaker_info = uploaded_speakers[voice_lower] - ref_text_from_upload = speaker_info.get("ref_text") - assert ref_text_from_upload == "Hi this is Alice" + speaker_info = uploaded_speakers[voice_lower] + ref_text_from_upload = speaker_info.get("ref_text") + assert ref_text_from_upload == "Hi this is Alice" - def test_uploaded_voice_without_ref_text_uses_request_ref_text(self): + def test_uploaded_voice_without_ref_text_uses_request_ref_text(self, mocker: MockerFixture): """If upload has no ref_text but request provides it, use request's.""" - request = MagicMock() + request = mocker.MagicMock() request.input = "Hello" request.voice = "bob" request.ref_audio = None diff --git a/tests/test_fish_speech_voice_cache.py b/tests/test_fish_speech_voice_cache.py index 8fe7a4a4d..1c299d801 100644 --- a/tests/test_fish_speech_voice_cache.py +++ b/tests/test_fish_speech_voice_cache.py @@ -10,11 +10,12 @@ Covers: import os import tempfile -from unittest.mock import MagicMock, patch +from pathlib import Path import numpy as np import pytest import torch +from pytest_mock import MockerFixture pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -61,18 +62,18 @@ class TestFishSpeechVoiceCacheIntegration: """Test the cache-hit / cache-miss / no-cache paths in the model.""" @pytest.fixture - def mock_model(self): + def mock_model(self, mocker: MockerFixture): """Create a mock FishSpeechSlowARForConditionalGeneration with cache.""" from vllm_omni.utils.voice_cache import VoiceEmbeddingCache - model = MagicMock() + model = mocker.MagicMock() model._voice_cache = VoiceEmbeddingCache(max_entries=4) model._semantic_begin_id = 151678 model._num_codebooks = 10 model._codebook_size = 4096 model.model_path = "/fake/model" - model.codebook_embeddings = MagicMock() - model.codebook_embeddings.weight = MagicMock() + model.codebook_embeddings = mocker.MagicMock() + model.codebook_embeddings.weight = mocker.MagicMock() model.codebook_embeddings.weight.device = torch.device("cpu") return model @@ -166,9 +167,13 @@ class TestFishSpeechVoiceCacheIntegration: class TestFishSpeechValidatorUploadedVoice: """Test _validate_fish_tts_request uploaded voice resolution.""" - def test_uploaded_voice_resolves_ref_audio(self): + def test_uploaded_voice_resolves_ref_audio( + self, + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, + ): """When voice matches an uploaded speaker, ref_audio should be auto-set.""" - request = MagicMock() + request = mocker.MagicMock() request.input = "Hello" request.voice = "alice" request.ref_audio = None @@ -185,17 +190,21 @@ class TestFishSpeechValidatorUploadedVoice: } # Simulate: voice in uploaded_speakers, file exists, get_audio returns data URL. - with patch("pathlib.Path.exists", return_value=True): - voice_lower = request.voice.lower() - assert voice_lower in uploaded_speakers + monkeypatch.setattr(Path, "exists", lambda self: True) - speaker_info = uploaded_speakers[voice_lower] - ref_text_from_upload = speaker_info.get("ref_text") - assert ref_text_from_upload == "Hi this is Alice" + voice_lower = request.voice.lower() + assert voice_lower in uploaded_speakers - def test_uploaded_voice_without_ref_text_uses_request_ref_text(self): + speaker_info = uploaded_speakers[voice_lower] + ref_text_from_upload = speaker_info.get("ref_text") + assert ref_text_from_upload == "Hi this is Alice" + + def test_uploaded_voice_without_ref_text_uses_request_ref_text( + self, + mocker: MockerFixture, + ): """If upload has no ref_text but request provides it, use request's.""" - request = MagicMock() + request = mocker.MagicMock() request.input = "Hello" request.voice = "bob" request.ref_audio = None