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

from __future__ import annotations

from contextlib import ExitStack

import pytest

from tests.e2e.features.helpers.custom_pipeline import (
    QwenImagePipelineWithLogProbForTest,
)
from tests.helpers.mark import hardware_test
from vllm_omni.entrypoints.async_omni import AsyncOmni

CUSTOM_PIPELINE_CLASS = QwenImagePipelineWithLogProbForTest
WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest"
MODEL = "tiny-random/Qwen-Image"

pytestmark = [pytest.mark.core_model]


@pytest.mark.diffusion
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.asyncio
async def test_collective_rpc_test_extension_name():
    with ExitStack() as after:
        engine = AsyncOmni(
            model=MODEL,
            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
            worker_extension_cls=WORKER_EXTENSION_CLASS,
            enforce_eager=True,
        )
        after.callback(engine.shutdown)

        result = await engine.collective_rpc(method="test_extension_name")

    assert isinstance(result, list), f"Expected list, got {type(result)}"
    assert result == [["vllm-omni-colocate-worker-extension-for-test"]]