"""Unit tests for format_response util."""


from abacus_contract.utils.format_response import \
    prepare_dataload_response, prepare_dataload_with_data_as_list_response


def test_prepare_dataload_response():
    """Test prepare_dataload_response function."""
    mock_data_1 = {
        'test_id': 1
    }
    mock_data_2 = {
        'test_id': 2
    }
    mock_list = [mock_data_1, mock_data_2]

    entity_ids = [1, 2, 3]

    res = prepare_dataload_response(
        entity_ids,
        mock_list,
        'test_id'
    )

    assert res == [
        {
            'data': mock_data_1
        },
        {
            'data': mock_data_2
        },
        {
            'data': None
        }
    ]


def test_prepare_dataload_with_data_as_list_response():
    """Test prepare_dataload_with_data_as_list_response function."""
    mock_data_1 = {
        'test_id': 1,
        'other': 1
    }
    mock_data_2 = {
        'test_id': 1,
        'other': 1
    }
    mock_data_3 = {
        'test_id': 2,
        'other': 1
    }
    mock_list = [mock_data_1, mock_data_2, mock_data_3]

    entity_ids = [1, 2, 3]

    res = prepare_dataload_with_data_as_list_response(
        entity_ids,
        mock_list,
        'test_id'
    )

    assert res == [
        {
            'data': [mock_data_1, mock_data_2]
        },
        {
            'data': [mock_data_3]
        },
        {
            'data': None
        }
    ]