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

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

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