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

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

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: