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

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

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: