import pytest

from vllm.model_executor.offloader.base import NoopOffloader
from vllm_ascend.model_executor.offloader import AscendPrefetchOffloader
from vllm_ascend.model_executor.offloader.selection import select_offload_layers
from vllm_ascend.worker.model_runner_v1 import create_ascend_weight_offloader

from tests.model_executor.offloader.test_runner_integration import make_config


@pytest.mark.parametrize(
    ("group_size", "expected_layers"),
    [
        (0, ()),
        (1, tuple(range(8))),
        (2, (1, 3, 5, 7)),
        (4, (3, 7)),
    ],
)
def test_offload_group_size_controls_layer_grouping(group_size, expected_layers):
    selection = select_offload_layers(
        num_layers=8,
        group_size=group_size,
        num_in_group=1,
    )

    assert selection.layer_indices == expected_layers
    assert selection.enabled is bool(expected_layers)


def test_offload_group_size_zero_keeps_default_noop_offloader():
    offloader = create_ascend_weight_offloader(make_config(0, 1))

    assert isinstance(offloader, NoopOffloader)


def test_positive_offload_group_size_enables_ascend_prefetch_offloader():
    offloader = create_ascend_weight_offloader(make_config(4, 1))

    assert isinstance(offloader, AscendPrefetchOffloader)
    assert offloader.group_size == 4
    assert offloader.num_in_group == 1
