"""Test for the Chart Handler.""" import datetime from unittest.mock import patch import pytest import uuid from charts.connectors import snowflake from charts.logic import sql_loader from charts.logic.dataloaders.charts.dates.placements import private_dataloader from charts.logic.permissions import AdminPermissionsType from charts.logic.permissions import NonAdminPermissionsType def test_dataloader_private_mismatched_row_count( mock_admin_permissions: AdminPermissionsType, ) -> None: """Test _dataloader_private when wrong row count returned.""" fetchall_mock_response = [ [ True, None ], [ True, "bad75dc4-8df3-45a5-baa5-ddb40dea3242,151a4684-fb9b-4ae9-b938-6da0f2ea1f10" ] ] with patch.object(snowflake, 'fetchall_nocache', return_value=fetchall_mock_response): with pytest.raises(Exception): private_dataloader._dataloader_private( [ { "chart_id": uuid.UUID("00637bf4-2eae-407c-a0b6-23aa85b6d47c"), "chart_date": datetime.date.fromisoformat("2020-01-01"), "epoch": 0, "uuid": uuid.UUID("aa8fc35c-665d-4b7b-b320-aa41161bba73"), }, ], show_sme_data=False, permissions=mock_admin_permissions, ) def test_dataloader_private_admin( mock_admin_permissions: AdminPermissionsType, ) -> None: """Test _dataloader_private for admin""" fetchall_mock_response = [ [ False, None, None, ], [ True, False, "", ], [ True, False, "8e25c4b9-ddf2-4609-a6bf-74ede9abddfe", ], [ True, True, "bad75dc4-8df3-45a5-baa5-ddb40dea3242,151a4684-fb9b-4ae9-b938-6da0f2ea1f10", ] ] with patch.object(private_dataloader, 'chart_placement_is_part_of_catalog_cte') as mock_is_part_of_catalog_sq,\ patch.object(sql_loader, 'load_query') as mock_sql_loader,\ patch.object(snowflake, 'fetchall_nocache') as mock_fetchall: mock_is_part_of_catalog_sq.return_value = "IS PART OF CATALOG SQL", { 'nfd_filter': [["'TestNFD1', 'TestNFD2'"]], } mock_sql_loader.return_value = 'MOCK SQL ({{ is_part_of_catalog_cte }})' mock_fetchall.return_value = fetchall_mock_response response = private_dataloader._dataloader_private( [ { "chart_id": uuid.UUID("00637bf4-2eae-407c-a0b6-23aa85b6d47c"), "chart_date": datetime.date.fromisoformat("2020-01-01"), "epoch": 0, "uuid": uuid.UUID("d1ed6afa-8d3b-473f-a419-2be350840000"), }, { "chart_id": uuid.UUID("00637bf4-2eae-407c-a0b6-23aa85b6d47c"), "chart_date": datetime.date.fromisoformat("2020-01-01"), "epoch": 0, "uuid": uuid.UUID("87c0522d-1f55-4267-901e-731b521c96b7"), }, { "chart_id": uuid.UUID("00637bf4-2eae-407c-a0b6-23aa85b6d47c"), "chart_date": datetime.date.fromisoformat("2020-01-01"), "epoch": 0, "uuid": uuid.UUID("9ecb8ef5-811d-4c7d-8e32-d1879a549c3e"), }, { "chart_id": uuid.UUID("00637bf4-2eae-407c-a0b6-23aa85b6d47c"), "chart_date": datetime.date.fromisoformat("2020-01-01"), "epoch": 0, "uuid": uuid.UUID("f666be83-d717-4080-8e49-36df81974039"), }, ], show_sme_data=True, permissions=mock_admin_permissions, ) assert response == [ { 'data': None, }, { 'data': { 'is_part_of_catalog': False, 'brand_ids': [], }, }, { 'data': { 'is_part_of_catalog': False, 'brand_ids': ['8e25c4b9-ddf2-4609-a6bf-74ede9abddfe'], }, }, { 'data': { 'is_part_of_catalog': True, 'brand_ids': ['bad75dc4-8df3-45a5-baa5-ddb40dea3242', '151a4684-fb9b-4ae9-b938-6da0f2ea1f10'], }, }, ] mock_sql_loader.assert_called_once_with('chart_placement/dataloader_private/admin') mock_fetchall.assert_called_once_with( 'MOCK SQL (IS PART OF CATALOG SQL)', { 'chart_placements': [ ["(0, '00637bf4-2eae-407c-a0b6-23aa85b6d47c', '2020-01-01', 'd1ed6afa-8d3b-473f-a419-2be350840000')"], ["(1, '00637bf4-2eae-407c-a0b6-23aa85b6d47c', '2020-01-01', '87c0522d-1f55-4267-901e-731b521c96b7')"], ["(2, '00637bf4-2eae-407c-a0b6-23aa85b6d47c', '2020-01-01', '9ecb8ef5-811d-4c7d-8e32-d1879a549c3e')"], ["(3, '00637bf4-2eae-407c-a0b6-23aa85b6d47c', '2020-01-01', 'f666be83-d717-4080-8e49-36df81974039')"], ], 'nfd_filter': [["'TestNFD1', 'TestNFD2'"]], } ) def test_dataloader_private_non_admin() -> None: """Test _dataloader_private for a non-admin.""" permissions: NonAdminPermissionsType = { 'artist_ids': [], 'vendor_ids': [123], 'subaccount_ids': [], 'label_participant_ids': [], } fetchall_mock_response = [ [ False, None, None, ], [ True, False, None, ], [ True, True, None, ] ] with patch.object(private_dataloader, 'chart_placement_is_part_of_catalog_cte') as mock_is_part_of_catalog_sq,\ patch.object(sql_loader, 'load_query') as mock_sql_loader,\ patch.object(snowflake, 'fetchall_nocache') as mock_fetchall: mock_is_part_of_catalog_sq.return_value = "QUERY", { 'nfd_filter': [["'TestNFD1', 'TestNFD2'"]], 'vendor_ids': [123], } mock_sql_loader.return_value = 'MOCK SQL {{ is_part_of_catalog_cte }}' mock_fetchall.return_value = fetchall_mock_response response = private_dataloader._dataloader_private( [ { "chart_id": uuid.UUID("00637bf4-2eae-407c-a0b6-23aa85b6d47c"), "chart_date": datetime.date.fromisoformat("2020-01-01"), "epoch": 0, "uuid": uuid.UUID("caaf8a53-e0ce-4144-ba92-07900c94bd45"), }, { "chart_id": uuid.UUID("00637bf4-2eae-407c-a0b6-23aa85b6d47c"), "chart_date": datetime.date.fromisoformat("2020-01-01"), "epoch": 0, "uuid": uuid.UUID("1c282c15-8da0-4d0d-912b-75f49fb5c512"), }, { "chart_id": uuid.UUID("00637bf4-2eae-407c-a0b6-23aa85b6d47c"), "chart_date": datetime.date.fromisoformat("2020-01-01"), "epoch": 0, "uuid": uuid.UUID("05aee4d8-a778-4d56-abb8-73369b81956d"), }, ], show_sme_data=False, permissions=permissions, ) assert response == [ { 'data': None, }, { 'data': { 'is_part_of_catalog': False, 'brand_ids': None, }, }, { 'data': { 'is_part_of_catalog': True, 'brand_ids': None, }, }, ] mock_sql_loader.assert_called_once_with('chart_placement/dataloader_private/user') mock_fetchall.assert_called_once_with( 'MOCK SQL QUERY', { 'chart_placements': [ ["(0, '00637bf4-2eae-407c-a0b6-23aa85b6d47c', '2020-01-01', 'caaf8a53-e0ce-4144-ba92-07900c94bd45')"], ["(1, '00637bf4-2eae-407c-a0b6-23aa85b6d47c', '2020-01-01', '1c282c15-8da0-4d0d-912b-75f49fb5c512')"], ["(2, '00637bf4-2eae-407c-a0b6-23aa85b6d47c', '2020-01-01', '05aee4d8-a778-4d56-abb8-73369b81956d')"], ], 'nfd_filter': [["'TestNFD1', 'TestNFD2'"]], 'vendor_ids': [123], } ) def test_dataloader_private_key_serializer_admin( mock_admin_permissions: AdminPermissionsType ) -> None: """Test _dataloader_private_key_serializer for an admin.""" result = private_dataloader._dataloader_private_key_serializer( { 'chart_id': uuid.UUID('0ef02e56-1e15-4557-ada5-d1d69cd5b4ca'), 'chart_date': datetime.date(2024, 1, 2), 'epoch': 0, 'uuid': uuid.UUID('c1468fdc-0a5d-445e-8af6-2433908ac470'), }, show_sme_data=False, permissions=mock_admin_permissions, ) assert result == 'chart:0ef02e56-1e15-4557-ada5-d1d69cd5b4ca:date:2024-01-02:epoch:0:placement:c1468fdc-0a5d-445e-8af6-2433908ac470:private:sme_data:False:admin' def test_dataloader_private_key_serializer_non_admin() -> None: """Test _dataloader_private_key_serializer for a non-admin.""" permissions: NonAdminPermissionsType = { 'vendor_ids': [], 'subaccount_ids': [1234], 'label_participant_ids': [5678, 9012], 'artist_ids': [], } result = private_dataloader._dataloader_private_key_serializer( { 'chart_id': uuid.UUID('0ef02e56-1e15-4557-ada5-d1d69cd5b4ca'), 'chart_date': datetime.date(2024, 1, 2), 'epoch': 0, 'uuid': uuid.UUID('c1468fdc-0a5d-445e-8af6-2433908ac470'), }, show_sme_data=False, permissions=permissions, ) assert result == "chart:0ef02e56-1e15-4557-ada5-d1d69cd5b4ca:date:2024-01-02:epoch:0:placement:c1468fdc-0a5d-445e-8af6-2433908ac470:private:sme_data:False:nonadmin:{'vendor_ids': [], 'subaccount_ids': [1234], 'label_participant_ids': [5678, 9012], 'artist_ids': []}"