import pytest

from vllm_ascend.model_executor.offloader.selection import (
    OffloadSelection,
    select_offload_layers,
)


def test_group_size_zero_disables_offload():
    selection = select_offload_layers(num_layers=8, group_size=0, num_in_group=1)
    assert selection == OffloadSelection(enabled=False, layer_indices=())


def test_num_in_group_zero_selects_no_layers():
    selection = select_offload_layers(num_layers=8, group_size=4, num_in_group=0)
    assert selection == OffloadSelection(enabled=False, layer_indices=())


def test_num_in_group_greater_than_group_size_raises():
    with pytest.raises(ValueError, match="offload_num_in_group .* must be <= offload_group_size"):
        select_offload_layers(num_layers=8, group_size=2, num_in_group=3)


def test_group_size_four_num_one_selects_last_layer_of_each_group():
    selection = select_offload_layers(num_layers=16, group_size=4, num_in_group=1)
    assert selection == OffloadSelection(enabled=True, layer_indices=(3, 7, 11, 15))


def test_group_size_eight_num_two_selects_last_two_layers_of_each_group():
    selection = select_offload_layers(num_layers=24, group_size=8, num_in_group=2)
    assert selection == OffloadSelection(enabled=True, layer_indices=(6, 7, 14, 15, 22, 23))


def test_num_in_group_equal_group_size_selects_all_layers():
    selection = select_offload_layers(num_layers=5, group_size=2, num_in_group=2)
    assert selection == OffloadSelection(enabled=True, layer_indices=(0, 1, 2, 3, 4))
