vllm-omni/tests/diffusion/test_diffusion_plugin_hooks.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""
Unit tests for diffusion engine plugin extensibility hooks.

This module tests:
- Platform hooks: get_diffusion_worker_cls, get_diffusion_model_runner_cls
- Registry API: register_diffusion_model
- Worker integration: model runner resolved via platform hook
"""

from types import SimpleNamespace
from unittest.mock import Mock, patch

import pytest

from vllm_omni.diffusion.registry import (
    _DIFFUSION_IR_OP_PRIORITY_FUNCS,
    _DIFFUSION_MODELS,
    _DIFFUSION_POST_PROCESS_FUNCS,
    _DIFFUSION_PRE_PROCESS_FUNCS,
    _DIFFUSION_PREFIX_CACHE_FUNCS,
    get_diffusion_post_process_func,
    get_diffusion_pre_process_func,
    get_diffusion_prefix_cache_func,
    register_diffusion_model,
)
from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


class TestPlatformDiffusionHooks:
    """Test OmniPlatform diffusion hook defaults."""

    def test_get_diffusion_worker_cls_default(self):
        """Test default diffusion worker class path."""
        result = OmniPlatform.get_diffusion_worker_cls()
        assert result == "vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker"

    def test_get_diffusion_model_runner_cls_default(self):
        """Test default diffusion model runner class path."""
        result = OmniPlatform.get_diffusion_model_runner_cls()
        assert result == "vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner"

    def test_get_diffusion_kv_block_tables_cls_default(self):
        """Test the default paged-KV BlockTables implementation."""
        from vllm.v1.worker.gpu.block_table import BlockTables

        assert OmniPlatform.get_diffusion_kv_block_tables_cls() is BlockTables

    def test_oot_enum_exists(self):
        """Test OOT is a valid platform enum value."""
        assert OmniPlatformEnum.OOT.value == "oot"

    def test_is_out_of_tree(self):
        """Test is_out_of_tree() returns True for OOT platform."""
        platform = OmniPlatform.__new__(OmniPlatform)
        platform._omni_enum = OmniPlatformEnum.OOT
        assert platform.is_out_of_tree() is True
        assert platform.is_cuda() is False


class TestRegisterDiffusionModel:
    """Test register_diffusion_model public API."""

    @pytest.fixture(autouse=True)
    def cleanup_registry(self):
        """Restore global registry dicts after each test."""
        original_models = _DIFFUSION_MODELS.copy()
        original_pre = _DIFFUSION_PRE_PROCESS_FUNCS.copy()
        original_post = _DIFFUSION_POST_PROCESS_FUNCS.copy()
        original_ir_op_priority = _DIFFUSION_IR_OP_PRIORITY_FUNCS.copy()
        original_prefix_cache = _DIFFUSION_PREFIX_CACHE_FUNCS.copy()
        yield
        _DIFFUSION_MODELS.clear()
        _DIFFUSION_MODELS.update(original_models)
        _DIFFUSION_PRE_PROCESS_FUNCS.clear()
        _DIFFUSION_PRE_PROCESS_FUNCS.update(original_pre)
        _DIFFUSION_POST_PROCESS_FUNCS.clear()
        _DIFFUSION_POST_PROCESS_FUNCS.update(original_post)
        _DIFFUSION_IR_OP_PRIORITY_FUNCS.clear()
        _DIFFUSION_IR_OP_PRIORITY_FUNCS.update(original_ir_op_priority)
        _DIFFUSION_PREFIX_CACHE_FUNCS.clear()
        _DIFFUSION_PREFIX_CACHE_FUNCS.update(original_prefix_cache)

    def test_register_new_model(self):
        """Test registering a new diffusion model with pre/post process functions."""
        register_diffusion_model(
            model_arch="TestPipeline",
            module_name="test_plugin.diffusion.pipeline",
            class_name="TestPipeline",
            pre_process_func_name="test_pre_process",
            post_process_func_name="test_post_process",
            ir_op_priority_func_name="test_ir_op_priority",
            prefix_cache_func_name="test_prefix_cache",
        )
        assert "TestPipeline" in _DIFFUSION_MODELS
        assert _DIFFUSION_MODELS["TestPipeline"] == (
            "test_plugin.diffusion.pipeline",
            "",
            "TestPipeline",
        )
        assert _DIFFUSION_PRE_PROCESS_FUNCS["TestPipeline"] == "test_pre_process"
        assert _DIFFUSION_POST_PROCESS_FUNCS["TestPipeline"] == "test_post_process"
        assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["TestPipeline"] == "test_ir_op_priority"
        assert _DIFFUSION_PREFIX_CACHE_FUNCS["TestPipeline"] == "test_prefix_cache"

    def test_registered_prefix_hook_is_loaded_from_plugin(self):
        hook = Mock()
        factory = Mock(return_value=hook)
        register_diffusion_model(
            "TestPrefixPipeline", "test_plugin.pipeline", "TestPrefixPipeline", prefix_cache_func_name="prefix_factory"
        )
        config = SimpleNamespace(model_class_name="TestPrefixPipeline")
        with patch(
            "vllm_omni.diffusion.registry.importlib.import_module", return_value=SimpleNamespace(prefix_factory=factory)
        ) as load:
            assert get_diffusion_prefix_cache_func(config) is hook
        load.assert_called_once_with("test_plugin.pipeline")
        factory.assert_called_once_with(config)

    def test_builtin_hunyuan_prefix_hook(self):
        from vllm_omni.diffusion.models.hunyuan_image3.request_layout import prepare_hunyuan_prefix_cache

        assert (
            get_diffusion_prefix_cache_func(SimpleNamespace(model_class_name="HunyuanImage3ForCausalMM"))
            is prepare_hunyuan_prefix_cache
        )

    def test_register_model_accepts_deprecated_action_postprocess_keyword(self):
        """Deprecated action hook keyword is accepted but not registered."""
        register_diffusion_model(
            model_arch="LegacyActionPipeline",
            module_name="test_plugin.diffusion.pipeline",
            class_name="LegacyActionPipeline",
            post_process_func_name="test_post_process",
            action_post_process_func_name="test_action_post_process",
        )

        assert "LegacyActionPipeline" in _DIFFUSION_MODELS
        assert _DIFFUSION_POST_PROCESS_FUNCS["LegacyActionPipeline"] == "test_post_process"


@pytest.mark.parametrize(
    ("model_class_name", "get_process_func"),
    [
        ("QwenImagePipeline", get_diffusion_post_process_func),
        ("WanImageToVideoPipeline", get_diffusion_post_process_func),
        ("WanImageToVideoPipeline", get_diffusion_pre_process_func),
    ],
)
def test_diffusers_backend_skips_native_process_hooks(model_class_name, get_process_func):
    od_config = SimpleNamespace(
        model_class_name=model_class_name,
        custom_pipeline_args=None,
        diffusion_load_format="diffusers",
    )

    with patch("vllm_omni.diffusion.registry._load_process_func") as load_process_func:
        assert get_process_func(od_config) is None

    load_process_func.assert_not_called()


@pytest.mark.parametrize(
    ("model_class_name", "get_process_func"),
    [
        ("QwenImagePipeline", get_diffusion_post_process_func),
        ("WanImageToVideoPipeline", get_diffusion_post_process_func),
        ("WanImageToVideoPipeline", get_diffusion_pre_process_func),
    ],
)
def test_custom_pipeline_preserves_native_process_hooks(model_class_name, get_process_func):
    process_func = Mock()
    od_config = SimpleNamespace(
        model_class_name=model_class_name,
        custom_pipeline_args={"pipeline_class": "custom.Pipeline"},
        diffusion_load_format="diffusers",
    )

    with patch("vllm_omni.diffusion.registry._load_process_func", return_value=process_func) as load_process_func:
        assert get_process_func(od_config) is process_func

    load_process_func.assert_called_once()


class TestWorkerUsesHook:
    """Test that DiffusionWorker resolves model runner via platform hook."""

    @patch("vllm_omni.diffusion.worker.diffusion_worker.resolve_obj_by_qualname")
    @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
    def test_model_runner_resolved_via_platform(self, mock_platform, mock_resolve):
        """Test model runner class is resolved from platform hook return value."""
        from unittest.mock import Mock

        from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker

        mock_runner_instance = Mock()
        mock_runner_cls = Mock(return_value=mock_runner_instance)
        mock_platform.get_diffusion_model_runner_cls.return_value = "custom.path"
        mock_resolve.return_value = mock_runner_cls

        with patch.object(DiffusionWorker, "init_device"):
            worker = DiffusionWorker(local_rank=0, rank=0, od_config=Mock(), skip_load_model=True)

        assert worker.model_runner is mock_runner_instance
        mock_platform.get_diffusion_model_runner_cls.assert_called_once()
        mock_resolve.assert_called_once_with("custom.path")

    @patch("vllm_omni.diffusion.worker.diffusion_worker.get_diffusion_ir_op_priority_func")
    @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
    def test_ir_op_priority_hook_receives_platform_default(self, mock_platform, mock_get_hook):
        """Test model IR priority hook merges from the platform default."""
        from vllm_omni.diffusion.worker.diffusion_worker import _resolve_ir_op_priority

        od_config = SimpleNamespace(model_class_name="TestPipeline")
        vllm_config = SimpleNamespace()
        default_priority = object()
        merged_priority = object()
        hook = Mock(return_value=merged_priority)
        mock_platform.get_default_ir_op_priority.return_value = default_priority
        mock_get_hook.return_value = hook

        assert _resolve_ir_op_priority(od_config, vllm_config) is merged_priority
        mock_platform.get_default_ir_op_priority.assert_called_once_with(vllm_config)
        mock_get_hook.assert_called_once_with(od_config)
        hook.assert_called_once_with(default_priority, vllm_config=vllm_config)