[Misc] Cleanup: use consistent pytest-mock in unit tests (#2698)

Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>

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
 

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"),

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

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)
 
 
 # ============================

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

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:

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"

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(

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)

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(

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,

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()

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

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(

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"
 

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,

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="<speaker:1>你好，欢迎回来。",
@@ -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

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

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."""

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(

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)

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]

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)

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

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