import pytest

from server.utils.parallel import DictResult, ListResult, request_in_chunks


@pytest.mark.parametrize(
    "chunk_size,chunks_count,result_type,ids_list,func_results,call_count,expected_result",
    (
        (1, 1, DictResult, list(range(2)), [{"f1": 1}, {"f2": 2}], 2, {"f1": 1, "f2": 2}),
        (2, 2, DictResult, list(range(3)), [{"a": 1, "b": "k"}, {"c": 3}], 2, {"a": 1, "b": "k", "c": 3}),
        (2, 2, DictResult, list(range(5)), [{"a": 1}, {"b": "k"}, {"c": 3}], 3, {"a": 1, "b": "k", "c": 3}),
        (1, 2, ListResult, list(range(3)), [[{"a": 1}, {"b": "k"}], [2], [5, 7]], 3, [{"a": 1}, {"b": "k"}, 2, 5, 7]),
    ),
)
@pytest.mark.noapp
async def test_request_in_chunks(
    chunk_size, chunks_count, result_type, ids_list, func_results, call_count, expected_result, mocker
):
    async def func(ids1, *args, **kwargs):
        return func_results[min(ids1) // chunk_size]

    mock_func = mocker.Mock(side_effect=func)

    decorated_func = request_in_chunks(
        chunk_size=chunk_size, chunks_count=chunks_count, items_key="ids1", result_type=result_type
    )(mock_func)

    result = await decorated_func(ids1=ids_list)

    assert result == expected_result
    assert mock_func.call_count == call_count