[Test] Cleanup for helpers (#6523)

Signed-off-by: wangyu <410167048@qq.com>
Signed-off-by: [Your Name] <your.email@example.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

tests/diffusion/cache/test_cache_dit.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
 Model specific tests for CacheDiT enablement.
@@ -13,6 +13,7 @@ from unittest.mock import Mock, patch
 import pytest
 import torch
 from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2
+from vllm.distributed import parallel_state
 
 import vllm_omni.diffusion.cache.cachedit as cd_backend
 import vllm_omni.diffusion.cache.cachedit.model_specific as cd_model_specific
@@ -61,6 +62,18 @@ def test_custom_cache_dit_enablers_are_registered_explicitly():
         assert cd_backend.CUSTOM_DIT_ENABLERS == expected_enablers
 
 
+@pytest.fixture()
+def init_fake_tp_group(mocker):
+    """Provide a fake TP group so vLLM linear layers can be instantiated."""
+    mock_tp = mocker.MagicMock()
+    mock_tp.world_size = 1
+    mock_tp.rank_in_group = 0
+    old = parallel_state._TP
+    parallel_state._TP = mock_tp
+    yield
+    parallel_state._TP = old
+
+
 def test_wan22_vace_uses_wan22_custom_cache_dit_enabler():
     assert cd_backend.CUSTOM_DIT_ENABLERS["Wan22VACEPipeline"] is cd_model_specific.enable_cache_for_wan22
 
@@ -110,7 +123,9 @@ def test_cachedit_consumers_use_package_api():
 
             for module in modules:
                 if module == legacy_module or module.startswith(internal_prefix):
-                    invalid_imports.append(f"{source_path.relative_to(package_root)}:{node.lineno}: {module}")
+                    invalid_imports.append(
+                        f"{source_path.relative_to(package_root)}:{getattr(node, 'lineno', '?')}: {module}"
+                    )
 
     assert not invalid_imports, "Cache-DiT consumers must use the package API:\n" + "\n".join(invalid_imports)
 

tests/diffusion/diffusion_backend/test_diffusers_backend.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from collections import namedtuple
 from types import SimpleNamespace
@@ -10,7 +10,7 @@ from diffusers import DiffusionPipeline  # pyright: ignore[reportPrivateImportUs
 from PIL import Image
 
 from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
 from vllm_omni.diffusion.data import (
     DiffusionOutput,
     DiffusionParallelConfig,
@@ -651,7 +651,7 @@ class TestDiffusersBackendEndToEndExecution:
     def test_t2i_random_weights(
         self,
         omni_server: OmniServer,
-        openai_client: OpenAIClientHandler,
+        online_client: OnlineOmniClient,
     ):
         messages = dummy_messages_from_mix_data(content_text="a photo of an astronaut riding a horse on mars")
 
@@ -667,8 +667,10 @@ class TestDiffusersBackendEndToEndExecution:
             },
         }
 
-        response = openai_client.send_diffusion_request(request_config)
-        image: Image.Image = response[0].images[0]  # pyright: ignore[reportOptionalSubscript]
+        response = online_client.send_diffusion_request(request_config)
+        images = response[0].images
+        assert images is not None
+        image: Image.Image = images[0]
 
         # Request config has incomplete width/height, so internal assertion in `send_diffusion_request` is incomplete.
         assert image.size == (512, 512)

tests/diffusion/distributed/test_tensor_parallel.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """DiT tensor-parallel parity and perf checks for diffusion models.
 
@@ -18,8 +18,8 @@ import pytest
 import torch
 from PIL import Image
 
-from tests.helpers.env import DeviceMemoryMonitor
 from tests.helpers.mark import hardware_test
+from tests.helpers.monitor import DeviceMemoryMonitor
 from tests.helpers.runtime import OmniRunner
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams

tests/diffusion/offloader/test_diffusion_cpu_offload.py

@@ -1,25 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import gc
+from typing import Any
 
 import numpy as np
 import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.helpers.env import DeviceMemoryMonitor
 from tests.helpers.mark import hardware_test
+from tests.helpers.monitor import DeviceMemoryMonitor
 from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
-AUDIO_MODEL = {
+AUDIO_MODEL: dict[str, dict[str, int | None]] = {
     "stabilityai/stable-audio-open-1.0": {"cuda": 100, "rocm": None},
 }
 
-IMAGE_MODELS = {
+IMAGE_MODELS: dict[str, dict[str, int | None]] = {
     "riverclouds/qwen_image_random": {"cuda": 2200, "rocm": 2100},
 }
 
-MODELS = {**AUDIO_MODEL, **IMAGE_MODELS}
+MODELS: dict[str, dict[str, int | None]] = {**AUDIO_MODEL, **IMAGE_MODELS}
 
 MODEL_MARKS = {
     "riverclouds/qwen_image_random": pytest.mark.core_model,
@@ -29,12 +33,12 @@ MODEL_MARKS = {
 _GATED_MODELS = {"stabilityai/stable-audio-open-1.0"}
 
 
-AUDIO_MODEL_PARAMS = {
+AUDIO_MODEL_PARAMS: dict[str, dict[str, Any]] = {
     "runner_params": {},
     "sampler_params": {},
 }
 
-IMAGE_MODEL_PARAMS = {
+IMAGE_MODEL_PARAMS: dict[str, dict[str, Any]] = {
     "runner_params": {},
     "sampler_params": {
         "height": 256,
@@ -43,7 +47,7 @@ IMAGE_MODEL_PARAMS = {
 }
 
 
-def inference(model_name: str, offload: bool = True):
+def inference(model_name: str, offload: bool = True) -> tuple[float, Any]:
     gc.collect()
     current_omni_platform.empty_cache()
     device_index = current_omni_platform.current_device()
@@ -133,6 +137,7 @@ def test_cpu_offload_diffusion_model(model_name: str):
     threshold = MODELS[model_name][platform]
     if threshold is None:
         pytest.skip(f"Threshold not defined for {platform} on {model_name}")
+    assert threshold is not None
 
     assert offload_peak_memory + threshold < no_offload_peak_memory, (
         f"Offload peak memory {offload_peak_memory} MB should be less than "

tests/diffusion/offloader/test_diffusion_layerwise_offload.py

@@ -1,36 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from typing import Any
+
 import numpy as np
 import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.helpers.env import DeviceMemoryMonitor
 from tests.helpers.mark import hardware_test
+from tests.helpers.monitor import DeviceMemoryMonitor
 from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
-AUDIO_MODEL = {
+AUDIO_MODEL: dict[str, dict[str, int | None]] = {
     "stabilityai/stable-audio-open-1.0": {"cuda": 1500, "rocm": 1500},
 }
 
-IMAGE_VIDEO_MODELS = {
+IMAGE_VIDEO_MODELS: dict[str, dict[str, int | None]] = {
     "riverclouds/qwen_image_random": {"cuda": 4500, "rocm": None},
     # "Wan-AI/Wan2.2-T2V-A14B-Diffusers": {"cuda": 45000, "rocm": None},
 }
 
-MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}
+MODELS: dict[str, dict[str, int | None]] = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}
 
 MODEL_MARKS = {
     "riverclouds/qwen_image_random": pytest.mark.core_model,
     "stabilityai/stable-audio-open-1.0": pytest.mark.full_model,
 }
 
-AUDIO_MODEL_PARAMS = {
+AUDIO_MODEL_PARAMS: dict[str, dict[str, Any]] = {
     "runner_params": {},
     "sampler_params": {},
 }
 
-IMAGE_VIDEO_MODELS_PARAMS = {
+IMAGE_VIDEO_MODELS_PARAMS: dict[str, dict[str, Any]] = {
     "runner_params": {"boundary_ratio": 0.875, "flow_shift": 5.0},
     "sampler_params": {"height": 480, "width": 640, "num_frames": 5},
 }
@@ -55,7 +60,7 @@ def run_inference(
     model_name: str,
     layerwise_offload: bool = False,
     num_inference_steps: int = 3,
-) -> float:
+) -> tuple[float, Any]:
     current_omni_platform.empty_cache()
     device_index = current_omni_platform.current_device()
     monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
@@ -140,6 +145,7 @@ def test_layerwise_offload_diffusion_model(model_name: str):
 
     if expected_saved_memory is None:
         pytest.skip(f"Threshold not defined for {platform} on {model_name}")
+    assert expected_saved_memory is not None
 
     # Verify that layerwise offloading significantly reduces memory usage
     # Passes only if the actual savings exceeds the expected savings

tests/diffusion/quantization/test_gguf_diffusion.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """End-to-end tests for GGUF quantization on diffusion models.
 
 Validates that GGUF-quantized diffusion models generate valid images and
@@ -19,8 +19,8 @@ import numpy as np
 import pytest
 import torch
 
-from tests.helpers.env import DeviceMemoryMonitor
 from tests.helpers.mark import hardware_test
+from tests.helpers.monitor import DeviceMemoryMonitor
 from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput