"""Tests for response utils."""

from royalties.utils import response


def test_prepare_dataload_response():
    """Test preparing dataload response."""
    entity_ids = [1, 2, 3, 4]
    entity_list = [
        {'test_id': 1, 'title': 'Test'},
        {'test_id': 2, 'title': 'Another Test'},
        {'test_id': 4, 'title': 'Yet More Test'},
    ]
    key_field_name = 'test_id'
    expected = [
        {'data': {'test_id': 1, 'title': 'Test'}},
        {'data': {'test_id': 2, 'title': 'Another Test'}},
        {'data': None},
        {'data': {'test_id': 4, 'title': 'Yet More Test'}},
    ]

    result = response.prepare_dataload_response(entity_ids, entity_list, key_field_name)

    assert result == expected