"""Unit tests for tasks for the Chartmetric Charts Ingestion Workflow.""" import datetime from datetime import timedelta from unittest.mock import call, MagicMock, patch from freezegun import freeze_time from kafka.structs import OffsetAndMetadata from kafka.structs import TopicPartition from oto import response import pytest from feed_ingestion.flows.chartmetric_charts import config from feed_ingestion.flows.chartmetric_charts import tasks @freeze_time('2020-01-01') @pytest.mark.parametrize( 'platform_names, in_expected', [ ( 'amazon', ['amazon'], ), ( 'amazon,appleMusic', ['amazon', 'appleMusic'], ), ( 'amazon, appleMusic', ValueError, ), ( 'amazn,appleMusic', ValueError, ), ( 'apple_music,amazon,deezer,linemusic,shazam,spotify' ',soundcloud,recochoku,tiktok,tiktok_weekly,youtube', ValueError, ), ( 'appleMusic,amazon,deezer,linemusic,shazam,spotify' ',soundcloud,recochoku,tiktok,tiktok_weekly,youtube', ValueError, ), ( 'appleMusic,amazon,deezer,linemusic,shazam' ',spotify,soundcloud,recochoku,tiktok,youtube,itunes', ['appleMusic', 'amazon', 'deezer', 'linemusic', 'shazam', 'spotify', 'soundcloud', 'recochoku', 'tiktok', 'youtube', 'itunes'], ), ]) def test_bootstrap_valid_platforms(monkeypatch, platform_names, in_expected): """Test how bootstrap handles platform_names passed.""" kwargs = dict( date=None, date_limit=None, platform_names=platform_names ) if in_expected is ValueError: with pytest.raises(in_expected): tasks.bootstrap(MagicMock(), **kwargs) else: result = tasks.bootstrap(MagicMock(), **kwargs) platform_days_back = {platform: ('2019-12-28' if platform == 'youtube' else '2019-12-30') for platform in in_expected} expected = { 'feed_name': 'chartmetric_charts', 'date': '2019-12-30', 'date_limit': '3000-01-01', 'ingestion_started_at': '2020-01-01 00:00:00', 'platform_names': in_expected, 'platform_days_back': platform_days_back, 'reload': 'False' } assert result == expected @freeze_time('2020-01-01') @pytest.mark.parametrize( 'days_back, expected_days_back', ( [None, '2019-12-30'], ['2', '2019-12-30'], ['1', '2019-12-31'], [1, '2019-12-31'], ) ) def test_bootstrap_default_date(monkeypatch, days_back, expected_days_back): """Test that bootstrap returns the default date.""" date = None date_limit = None days_back_yt = int(days_back) + 2 if days_back else 4 expected_days_back_yt = (datetime.datetime.today() - datetime.timedelta( days=days_back_yt)).date().strftime('%Y-%m-%d') expected_platform_days_back = {platform: ( expected_days_back_yt if platform == 'youtube' else expected_days_back) for platform in config.all_platform_names} expected = { 'feed_name': 'chartmetric_charts', 'date': expected_days_back, 'date_limit': '3000-01-01', 'ingestion_started_at': '2020-01-01 00:00:00', 'platform_names': config.all_platform_names, 'platform_days_back': expected_platform_days_back, 'reload': 'False', } result = tasks.bootstrap( MagicMock(), date, date_limit, days_back=days_back) assert result == expected @freeze_time('2020-01-01') def test_bootstrap_context_date(monkeypatch): """Test that bootstrap returns the context date.""" date = '2019-01-01' date_limit = '2019-01-02' date_obj = datetime.datetime.strptime(date, '%Y-%m-%d') date_obj_yt = date_obj - timedelta(days=2) expected_platform_days_back = {platform: (date_obj_yt.strftime('%Y-%m-%d') if platform == 'youtube' else date_obj.strftime('%Y-%m-%d')) for platform in config.all_platform_names} expected = { 'feed_name': 'chartmetric_charts', 'date': date, 'date_limit': date_limit, 'ingestion_started_at': '2020-01-01 00:00:00', 'platform_names': config.all_platform_names, 'platform_days_back': expected_platform_days_back, 'reload': 'False', } result = tasks.bootstrap(MagicMock(), date, date_limit) assert result == expected @patch('feed_ingestion.flows.chartmetric_charts.tasks.get_sf_config') @patch('feed_ingestion.flows.chartmetric_charts.tasks.SnowflakeExecutor') def test_clear_log_table(mock_executor_class, mock_get_sf_config): """Test clear log table.""" mock_conifg = {} mock_get_sf_config.return_value = mock_conifg mock_executor = MagicMock() mock_executor_class.return_value.__enter__.return_value = mock_executor feed_name = 'chartmetric_charts' date = '2020-01-01' date_limit = '3000-01-01' platform_name = 'spotify' force_refresh = 'True' tasks.clear_log_table( MagicMock(), feed_name, date, date_limit, platform_name, force_refresh) mock_get_sf_config.assert_called_once_with(feed_name) mock_executor_class.assert_called_once_with(mock_conifg) mock_executor.execute_query.assert_called_once_with( 'clear_log_table', log_table_name=config.log_table_name, ingest_date=date, ingest_date_limit=date_limit, platform_name=platform_name) mock_executor.execute_query.reset_mock() force_refresh = 'False' tasks.clear_log_table( MagicMock(), feed_name, date, date_limit, platform_name, force_refresh) assert mock_executor.execute_query.call_count == 0 @patch('feed_ingestion.flows.chartmetric_charts.tasks.get_sf_config') @patch('feed_ingestion.flows.chartmetric_charts.tasks.SnowflakeExecutor') def test_insert_into_log_table(mock_executor_class, mock_get_sf_config): """Test insert_into_log_table.""" mock_conifg = {} mock_get_sf_config.return_value = mock_conifg mock_executor = MagicMock() mock_executor_class.return_value.__enter__.return_value = mock_executor feed_name = 'chartmetric_charts' date = '2020-01-01' date_limit = '3000-01-01' ingestion_started_at = '2020-01-01 00:00:00' tasks.insert_into_log_table( MagicMock(), feed_name, date, date_limit, ingestion_started_at) mock_get_sf_config.assert_called_once_with(feed_name) mock_executor_class.assert_called_once_with(mock_conifg) mock_executor.execute_query.assert_called_once_with( 'insert_into_log_table', ingest_date='2020-01-01', ingest_date_limit='3000-01-01', ingestion_started_at='2020-01-01 00:00:00', log_table_name='chartmetric_charts_log', table_name='STAGING_FACT_CHARTMETRIC_CHARTS' ) @patch('feed_ingestion.flows.chartmetric_charts.tasks.get_sf_config') @patch('feed_ingestion.flows.chartmetric_charts.tasks.SnowflakeExecutor') def test_create_staging_fact(mock_executor_class, mock_get_sf_config): """Test create_staging_fact.""" mock_executor = MagicMock() mock_executor_class.return_value.__enter__.return_value = mock_executor tasks.create_staging_fact(MagicMock()) kwargs = {'table_name': 'STAGING_FACT_CHARTMETRIC_CHARTS'} mock_executor.execute_query.assert_called_once_with( 'create_staging_fact', **kwargs ) mock_executor.\ grant_select_to_facts_db_prod_schema_read.assert_called_once_with( 'STAGING_FACT_CHARTMETRIC_CHARTS') @pytest.mark.parametrize('platform_name, query_name', [ ('amazon', 'load_staging_fact_amazon'), ('appleMusic', 'load_staging_fact_applemusic'), ('spotify', 'load_staging_fact_spotify'), ('tiktok', 'load_staging_fact_tiktok'), ('itunes', 'load_staging_fact_itunes'), ]) @patch('feed_ingestion.flows.chartmetric_charts.tasks.get_sf_config') @patch('feed_ingestion.flows.chartmetric_charts.tasks.SnowflakeExecutor') def test_load_staging_fact(mock_executor_class, mock_get_sf_config, platform_name, query_name): """Test load_staging_fact.""" mock_executor = MagicMock() mock_executor_class.return_value.__enter__.return_value = mock_executor date = '2020-01-01' date_limit = '3000-01-01' tasks.load_staging_fact(MagicMock(), date, date_limit, platform_name) if platform_name is None: mock_executor.execute_query.assert_not_called() else: params_expected = { 'platform_name': platform_name, 'ingest_date': date, 'ingest_date_limit': date_limit } mock_executor.execute_query.assert_called_once_with( query_name, **params_expected ) @patch('feed_ingestion.flows.chartmetric_charts.tasks.get_sf_config') @patch('feed_ingestion.flows.chartmetric_charts.tasks.SnowflakeExecutor') def test_update_dim_tables(mock_executor_class, mock_get_sf_config): """Test update_dim_tables.""" mock_executor = MagicMock() mock_executor_class.return_value.__enter__.return_value = mock_executor tasks.update_dim_tables(MagicMock()) mock_executor.execute_query.assert_called_once_with( 'update_dim_chart', ) @patch('feed_ingestion.flows.chartmetric_charts.tasks.get_sf_config') @patch('feed_ingestion.flows.chartmetric_charts.tasks.SnowflakeExecutor') @patch('feed_ingestion.flows.chartmetric_charts.tasks.time') def test_update_staging_fact_sound_recording( mock_time, mock_executor_class, mock_get_sf_config): """Test update_staging_fact_sound_recording.""" mock_executor = MagicMock() mock_executor_class.return_value.__enter__.return_value = mock_executor tasks.update_staging_fact_sound_recording(MagicMock()) mock_executor.execute_query.assert_has_calls([ call('update_staging_fact_sound_recording_by_isrc'), call('update_staging_fact_sound_recording_by_track_id'), ]) @patch('feed_ingestion.flows.chartmetric_charts.tasks.get_sf_config') @patch('feed_ingestion.flows.chartmetric_charts.tasks.SnowflakeExecutor') def test_update_staging_fact_public_product( mock_executor_class, mock_get_sf_config): """Test update_staging_fact_public_product.""" mock_executor = MagicMock() mock_executor_class.return_value.__enter__.return_value = mock_executor tasks.update_staging_fact_public_product(MagicMock()) mock_executor.execute_query.assert_called_once_with( 'update_staging_fact_public_product', ) @patch('feed_ingestion.flows.chartmetric_charts.tasks.get_sf_config') @patch('feed_ingestion.flows.chartmetric_charts.tasks.SnowflakeExecutor') def test_update_isrc_to_label_participant_mapping( mock_executor_class, mock_get_sf_config): """Test update_isrc_to_label_participant_mapping.""" mock_executor = MagicMock() mock_executor_class.return_value.__enter__.return_value = mock_executor tasks.update_isrc_to_label_participant_mapping(MagicMock()) mock_executor.execute_query.assert_called_once_with( 'update_isrc_to_label_participant_mapping', ) @patch('feed_ingestion.flows.chartmetric_charts.tasks.get_sf_config') @patch('feed_ingestion.flows.chartmetric_charts.tasks.SnowflakeExecutor') def test_update_upc_to_label_participant_mapping( mock_executor_class, mock_get_sf_config): """Test update_upc_to_label_participant_mapping.""" mock_executor = MagicMock() mock_executor_class.return_value.__enter__.return_value = mock_executor tasks.update_upc_to_label_participant_mapping(MagicMock()) mock_executor.execute_query.assert_called_once_with( 'update_upc_to_label_participant_mapping', ) @patch('feed_ingestion.flows.chartmetric_charts.tasks.get_sf_config') @patch('feed_ingestion.flows.chartmetric_charts.tasks.SnowflakeExecutor') def test_update_dim_chart_latest_chart_date( mock_executor_class, mock_get_sf_config): """Test update_dim_tables.""" mock_executor = MagicMock() mock_executor_class.return_value.__enter__.return_value = mock_executor tasks.update_dim_chart_latest_chart_date(MagicMock()) mock_executor.execute_query.assert_called_once_with( 'update_dim_chart_latest_chart_date', ) @pytest.mark.parametrize('platform_name', [ ('appleMusic'), ('spotify'), ('tiktok'), ('itunes'), ]) @patch('feed_ingestion.flows.chartmetric_charts.tasks.get_sf_config') @patch('feed_ingestion.flows.chartmetric_charts.tasks.SnowflakeExecutor') def test_load_fact_data(mock_executor_class, mock_get_sf_config, platform_name): """Test load_fact_data.""" mock_executor = MagicMock() mock_executor_class.return_value.__enter__.return_value = mock_executor date = '2020-01-01' date_limit = '3000-01-01' tasks.load_fact_data(MagicMock(), date, date_limit, platform_name) if platform_name is None: mock_executor.execute_query.assert_not_called() else: params_expected = { 'platform_name': platform_name, 'ingest_date': date, 'ingest_date_limit': date_limit } mock_executor.execute_query.assert_has_calls([ call('delete_from_fact_charts', **params_expected), call('load_fact_charts', **params_expected), ]) @pytest.mark.parametrize( 'platform_name, query_name, topic_suffix', [ ('amazon', 'get_sf_globalsoundrecording_relations', 'globalsoundrecording.relations'), ('appleMusic', 'get_sf_globalsoundrecording_relations', 'globalsoundrecording.relations'), ('spotify', 'get_sf_soundrecording_noisrc_nodes', 'soundrecording.noisrc.nodes'), ('tiktok', 'get_sf_soundrecording_noisrc_nodes', 'soundrecording.noisrc.nodes'), ('recochoku', 'get_sf_trackandparticipant_relations', '{}.trackandparticipant_relations'), ('linemusic', 'get_sf_trackandparticipant_relations', '{}.trackandparticipant_relations'), ('deezer', 'get_sf_productandlabel_relations', '{}.productandlabel.relations'), ('soundcloud', 'get_sf_productandlabel_relations', '{}.productandlabel.relations'), ] ) @patch('feed_ingestion.flows.chartmetric_charts.tasks.uuid4') @patch('feed_ingestion.flows.chartmetric_charts.tasks.get_sf_config') @patch('feed_ingestion.flows.chartmetric_charts.tasks.SnowflakeExecutor') @patch('feed_ingestion.flows.chartmetric_charts.tasks.KafkaExecutor') def test_push_metadata_to_kafka( mock_kafka_executor, mock_sf_executor_class, mock_get_sf_config, mock_uuid, platform_name, query_name, topic_suffix): """Test load_fact_data.""" mock_uuid.return_value = 'test_uuid' mock_sf_executor = MagicMock() mock_sf_config = {'db': 'DB', 'schema': 'SCHEMA'} mock_rows = [{'key': 'value'} for _ in range(1000)] mock_get_sf_config.return_value = mock_sf_config mock_sf_executor.fetchmany_dict_query.side_effect = ((mock_rows, []), []) (mock_sf_executor_class .return_value .__enter__ .return_value) = mock_sf_executor mock_producer = mock_kafka_executor().producer() topic_template = f'{config.KAFKA_TOPIC_PREFIX}{topic_suffix}' topic = topic_template.format(str(platform_name).lower()) return_value = tasks.push_metadata_to_kafka( MagicMock(), run_id='test_run_id', sf_query_name=query_name, kafka_topic_suffix=topic_suffix, message_key='key', platforms_to_skip=[], platform_name=platform_name) msg_key = { 'key': str(mock_rows[0].get('key')), 'SF_QUERY_ID': 'test_uuid' } if query_name == 'get_sf_productandlabel_relations': album_id_column = config.ALBUM_ID_COLUMN_MAPPING.get(platform_name) else: album_id_column = None params_expected_start = { 'table_name': 'STAGING_FACT_CHARTMETRIC_CHARTS', 'platform': platform_name, 'album_id_column': album_id_column } assert mock_sf_executor.fetchmany_dict_query.call_args_list == [ call(query_name=query_name, size=config.sf_batch_size, **params_expected_start), ] mock_producer.send.assert_called_with( topic, key=msg_key, value=mock_rows[0]) assert mock_producer.flush.called assert isinstance(return_value, dict) assert return_value['sf_query_id'] == 'test_uuid' @pytest.mark.parametrize( 'platform_name, topic_suffix, skip_platform, latest_offsets', [ ('amazon', 'globalsoundrecording.relations', False, {0: 11, 1: 45, 2: 88}), # noqa ('appleMusic', 'globalsoundrecording.relations', False, {0: 11, 1: 45, 2: 88}), # noqa ('spotify', 'soundrecording.noisrc.nodes', False, {0: 11, 1: 45, 2: 88}), # noqa ('tiktok', 'soundrecording.noisrc.nodes', False, {0: 11, 1: 45, 2: 88}), # noqa ('recochoku', '{}.trackandparticipant_relations', False, {0: 11, 1: 45, 2: 88}), # noqa ('linemusic', '{}.trackandparticipant_relations', False, {0: 11, 1: 45, 2: 88}), # noqa ('deezer', '{}.productandlabel.relations', False, {0: 11, 1: 45, 2: 88}), # noqa ('soundcloud', '{}.productandlabel.relations', True, {0: 11, 1: 45, 2: 88}), # noqa ] ) @patch('feed_ingestion.flows.chartmetric_charts.tasks.KafkaExecutor') def test_check_completion_status_in_kafka( mock_kafka_executor, platform_name, topic_suffix, skip_platform, latest_offsets): """Test check_completion_status_in_kafka.""" topic_template = f'{config.KAFKA_TOPIC_PREFIX}{topic_suffix}' topic = topic_template.format(str(platform_name).lower()) consumer_offsets = { TopicPartition(topic, 0): OffsetAndMetadata(12, 'test'), TopicPartition(topic, 1): OffsetAndMetadata(46, 'test'), TopicPartition(topic, 2): OffsetAndMetadata(89, 'test'), } client = mock_kafka_executor().admin_client() client.list_consumer_group_offsets = MagicMock( return_value=consumer_offsets) tasks.check_completion_status_in_kafka( MagicMock(), run_id='test_run_id', kafka_topic_suffix=topic_suffix, latest_offsets=latest_offsets, platform_name=platform_name, skip_kafka_step=skip_platform) if skip_platform: assert not client.list_consumer_group_offsets.called else: assert client.list_consumer_group_offsets.called @pytest.mark.parametrize('skip_platform', [False, True]) @patch('feed_ingestion.flows.chartmetric_charts.tasks.KafkaExecutor') def test_check_for_errors_in_dlq(mock_kafka_executor, skip_platform): """Test check_completion_status_in_kafka.""" consumer = mock_kafka_executor().consumer() tasks.check_for_errors_in_dlq( MagicMock(), run_id='test_run_id', workflow_id='test_workflow_id', sf_query_id='test_uuid', skip_kafka_step=skip_platform) call_list = [ bool(consumer.partitions_for_topic.called), bool(consumer.end_offsets.called), bool(consumer.assign.called), bool(consumer.seek_to_beginning.called), bool(consumer.poll.called) ] if skip_platform: assert not any(call_list) else: assert all(call_list) @pytest.mark.parametrize('platform_name, query_name', [ ('amazon', 'find_unchanged_charts_amazon'), ('appleMusic', 'find_unchanged_charts_applemusic'), ('spotify', 'find_unchanged_charts_spotify'), ('tiktok', 'find_unchanged_charts_tiktok'), ('itunes', 'find_unchanged_charts_itunes'), ]) @patch('feed_ingestion.flows.chartmetric_charts.tasks.get_sf_config') @patch('feed_ingestion.flows.chartmetric_charts.tasks.SnowflakeExecutor') def test_find_unchanged_charts(mock_executor_class, mock_get_sf_config, platform_name, query_name): """Test find_unchanged_charts.""" mock_executor = MagicMock() mock_executor_class.return_value.__enter__.return_value = mock_executor date = '2020-01-01' date_limit = '2020-01-04' tasks.find_unchanged_charts(MagicMock(), date, date_limit, platform_name) params_expected = { 'platform_name': platform_name, 'ingest_date': date, 'ingest_date_limit': date_limit } mock_executor.fetchall_query.assert_called_once_with( query_name, **params_expected ) @patch( 'feed_ingestion.flows.chartmetric_charts.config.OWS_CHARTS_BATCH_SIZE', 5, ) @patch('feed_ingestion.flows.chartmetric_charts.tasks.prime_chart_dates') @patch('feed_ingestion.flows.chartmetric_charts.tasks.get_sf_config') @patch('feed_ingestion.flows.chartmetric_charts.tasks.SnowflakeExecutor') def test_ows_charts_cache_prime( mock_executor_class, mock_get_sf_conf, prime_chart_dates, ): """Test ows_charts_cache_prime success.""" mock_executor = MagicMock() mock_executor_class.return_value.__enter__.return_value = mock_executor def mock_fetchall_query(_query_name, platform, date): if platform == 'youtube': return ( ( '37432e4e-0b61-46d5-8a13-64e448f766e9', datetime.date.fromisoformat('2020-01-04'), ), ( '0ae7e186-6fa3-49cf-9e73-f7ee3c18a4e5', datetime.date.fromisoformat('2020-01-02'), ), ) if platform == 'appleMusic': return ( ( 'af09deb0-8de0-4497-8f8a-825369519c0a', datetime.date.fromisoformat('2020-01-04'), ), ( '11c67436-6180-43f3-ab33-00c52b53b102', datetime.date.fromisoformat('2020-01-02'), ), ) mock_executor.fetchall_query.side_effect = mock_fetchall_query prime_chart_dates.return_value = response.Response({'ok': True}) date_limit = '2020-01-04' platform_days_back = { 'appleMusic': '2020-01-02', 'youtube': '2020-01-01', } tasks.ows_charts_cache_prime( MagicMock(), date_limit, platform_days_back, ) expected_apple_params = { 'platform': 'appleMusic', 'date': '2020-01-02', } expected_youtube_params = { 'platform': 'youtube', 'date': '2020-01-01', } mock_executor.fetchall_query.assert_has_calls([ call( 'get_platform_chart_ids', **expected_apple_params, ), call( 'get_platform_chart_ids', **expected_youtube_params, ), ]) prime_chart_dates.assert_has_calls([ call([ { 'chart_id': 'af09deb0-8de0-4497-8f8a-825369519c0a', 'chart_date': '2020-01-02', }, { 'chart_id': 'af09deb0-8de0-4497-8f8a-825369519c0a', 'chart_date': '2020-01-03', }, { 'chart_id': '11c67436-6180-43f3-ab33-00c52b53b102', 'chart_date': '2020-01-02', }, { 'chart_id': '37432e4e-0b61-46d5-8a13-64e448f766e9', 'chart_date': '2020-01-01', }, { 'chart_id': '37432e4e-0b61-46d5-8a13-64e448f766e9', 'chart_date': '2020-01-02', }, ]), call([ { 'chart_id': '37432e4e-0b61-46d5-8a13-64e448f766e9', 'chart_date': '2020-01-03', }, { 'chart_id': '0ae7e186-6fa3-49cf-9e73-f7ee3c18a4e5', 'chart_date': '2020-01-01', }, { 'chart_id': '0ae7e186-6fa3-49cf-9e73-f7ee3c18a4e5', 'chart_date': '2020-01-02', }, ]), ]) @patch('feed_ingestion.flows.chartmetric_charts.tasks.prime_chart_dates') @patch('feed_ingestion.flows.chartmetric_charts.tasks.get_sf_config') @patch('feed_ingestion.flows.chartmetric_charts.tasks.SnowflakeExecutor') def test_ows_charts_cache_prime_no_charts( mock_executor_class, mock_get_sf_conf, prime_chart_dates, ): """Test ows_charts_cache_prime success with no charts to prime.""" mock_executor = MagicMock() mock_executor_class.return_value.__enter__.return_value = mock_executor mock_executor.fetchall_query.return_value = tuple() date_limit = '2020-01-03' platform_days_back = { 'youtube': '2020-01-01', } tasks.ows_charts_cache_prime( MagicMock(), date_limit, platform_days_back, ) expected_params = { 'platform': 'youtube', 'date': '2020-01-01', } mock_executor.fetchall_query.assert_called_once_with( 'get_platform_chart_ids', **expected_params, ) prime_chart_dates.assert_not_called()