"""Test format response util."""

from payment.utils.format_response import (
    prepare_dataload_response,
    prepare_dataload_with_data_as_list_response,
)


def test_prepare_dataload_response():
    """Test prepare_dataload_response."""
    entity_ids = [1, 2, 3]
    entity_list = [{'id': 1}, {'id': 2}]
    key_field_name = 'id'

    expected_result = [{'data': {'id': 1}}, {'data': {'id': 2}}, {'data': None}]

    result = prepare_dataload_response(entity_ids, entity_list, key_field_name)
    assert result == expected_result


def test_prepare_dataload_with_data_as_list_response():
    """Test validation error creates an error response."""
    entity_ids = [1, 2, 3]
    entity_list = [{'id': 1}, {'id': 2}]
    key_field_name = 'id'

    expected_result = [{'data': [{'id': 1}]}, {'data': [{'id': 2}]}, {'data': None}]

    result = prepare_dataload_with_data_as_list_response(
        entity_ids, entity_list, key_field_name
    )
    assert result == expected_result