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

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

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