"""Unit tests for tasks for ingest iTunes raw data into staging_raw_itunes.""" from collections import OrderedDict from unittest import mock from unittest.mock import call from unittest.mock import MagicMock from unittest.mock import patch import boto3 from garcon import activity from garcon_contrib.dynamo_feed_status import \ garcon_feed_status import pytest from feed_ingestion.conf.config import SF_CONFIG from feed_ingestion.flows import helpers from feed_ingestion.flows.itunes import load_raw_table_tasks, tasks from feed_ingestion.flows.itunes.util import itunes_helpers from feed_ingestion.util import task_status from feed_ingestion.util.aws import s3 as s3utils MOCK_AWS_CONFIG = { 'access_key': 'access_key', 'access_secret': 'access_secret' } MOCK_SF_CONFIG = { 'db': 'db', 'schema': 'schema', 'role': 'role', 'warehouse': 'warehouse', } @pytest.fixture def sql_loader_version(): """Mock sql_loader instance.""" path = ('feed_ingestion.flows.itunes.' 'load_raw_table_tasks.sql_loader_version') with patch(path) as loader: loader.return_value = MagicMock() yield loader @pytest.fixture def mock_sql_loader(): """Mock sql_loader instance.""" with patch('feed_ingestion.flows.itunes.' 'load_raw_table_tasks.sql_loader') as loader: yield loader @pytest.fixture def mock_load_temp_staging_raw_table(): """Mock sql_loader instance.""" path = ('feed_ingestion.flows.itunes.' 'load_raw_table_tasks.load_temp_staging_raw_table') with patch(path) as mock_load_temp: yield mock_load_temp @pytest.fixture def mock_config(): """Mock config.""" with patch('feed_ingestion.flows.itunes.' 'load_raw_table_tasks.config') as mconfig: mconfig.feed_name = 'iTunes' mconfig.accounts = OrderedDict([ ('ioda', { 'destination_s3_path': ( 's3://bucket/iTunes/archives/{date:%Y-%m-%d}/{file_name}'), 'vendor_id': 80029727, 'file_info': { 'create_query_name': 'create_temp_table', 'load_query_name': 'load_temp_table', 'date_type': 'Daily', 'file_name': 'D_D_{vendor_id}_{date:%Y%m%d}.txt.gz', 'temp_table': 'staging_itunes_ioda_{date:%Y%m%d}', 'report_subtype': 'Detailed', 'report_type': 'Sales', }, }), ('bfm_orchard', { 'destination_s3_path': ( 's3://bucket/iTunes/archives/{date:%Y-%m-%d}/{file_name}'), 'vendor_id': 85420853, 'file_info': { 'create_query_name': 'create_temp_table', 'load_query_name': 'load_temp_table', 'date_type': 'Daily', 'file_name': 'D_D_{vendor_id}_{date:%Y%m%d}.txt.gz', 'temp_table': 'staging_itunes_bfm_{date:%Y%m%d}', 'report_subtype': 'Detailed', 'report_type': 'Sales', }, }), ('orchard', { 'destination_s3_path': ( 's3://bucket/iTunes/archives/{date:%Y-%m-%d}/{file_name}'), 'vendor_id': 80028967, 'file_info': { 'create_query_name': 'create_temp_table', 'load_query_name': 'load_temp_table', 'date_type': 'Daily', 'file_name': 'D_D_{vendor_id}_{date:%Y%m%d}.txt.gz', 'temp_table': 'staging_itunes_orchard_{date:%Y%m%d}', 'report_subtype': 'Detailed', 'report_type': 'Sales' }, }), ]) yield mconfig @pytest.fixture def mock_sf_config(): """Mock Snowflake config.""" with patch.dict(SF_CONFIG, MOCK_SF_CONFIG, clean=True) as mconfig: yield mconfig @pytest.fixture def mock_executor_context(sf_config_mock): """Yield executor context.""" sf_executor_class_path = ( 'feed_ingestion.flows.itunes.' 'load_raw_table_tasks.SnowflakeSQLExecutor') with patch(sf_executor_class_path) as sf_executor: mock_executor_context = \ sf_executor.return_value.__enter__.return_value yield mock_executor_context @pytest.fixture def mock_task_status(): """Yield task status.""" task_status_path = ( 'feed_ingestion.flows.itunes.load_raw_table_tasks.task_status') with patch(task_status_path) as task_status: task_status.is_completed_task.return_value = False task_status.mark_completed_task = MagicMock() yield task_status def test_bootstrap(mock_config, monkeypatch): """Test bootstrap.""" mock_config.s3 = { 'archive_bucket': 's3://test_s3_bucket/itunes/archives/{{date}}/'} orchard_vendors = ['80028547', '80028967', '80029727', '80031727', '80031813', '80032046', '80032084', '80032543', '80034465', '80035150', '85420853', '85320019'] mock_config.licensors = {'theorchard': orchard_vendors} monkeypatch.setattr( garcon_feed_status, 'get_overall_status', MagicMock( return_value='NOT_AVAILABLE')) monkeypatch.setattr( task_status, 'is_completed_overall_job', MagicMock( return_value=False)) monkeypatch.setattr( garcon_feed_status, 'set_missing_files', MagicMock( return_value=None)) monkeypatch.setattr( task_status, 'delete_newcontexts', MagicMock( return_value=None)) monkeypatch.setattr( task_status, 'get_report_contexts', MagicMock( return_value=True)) resp = tasks.bootstrap(MagicMock(), '2020-04-20', None, 'theorchard') resp_vendors = resp['reporter_accounts'].keys() assert resp['date'] == '2020-04-20' assert set(resp_vendors) == set(orchard_vendors) assert resp['s3_archive_bucket'] == 's3://dev-cucumbers/iTunes/' \ 'archives/2020-04-20/' assert resp['snowflake_error_limit'] == 1 def test_populate_temp_table( mock_config, sf_config_mock, mock_executor_context, sql_loader_version, monkeypatch, mock_load_temp_staging_raw_table, mock_task_status): """Test populate_temp_table task.""" mock_activity = MagicMock() monkeypatch.setattr( helpers, 'check_s3_key_exist', MagicMock(return_value=True)) load_raw_table_tasks.populate_temp_tables( mock_activity, '2015-10-09', 'feed_name', MOCK_AWS_CONFIG, sf_config_mock, {'vendor_id': '80029727', 'temp_table': 'staging_raw_itunes_theorchard_80029727_20151009', 'destination_s3_path': 's3://bucket/iTunes/archives/2015-10-09/' 'D_D_80029727_20151009.txt.gz'}, snowflake_error_limit=1) mock_executor_context.execute_query.assert_called_once() kwargs = {'query_name': 'load_temp_table', 'error_on_column_count_mismatch': 'True'} mock_load_temp_staging_raw_table.assert_has_calls([ call( mock_activity, '2015-10-09', 'iTunes', MOCK_AWS_CONFIG, 's3://bucket/iTunes/archives/2015-10-09/' 'D_D_80029727_20151009.txt.gz', 'staging_raw_itunes_theorchard_80029727_20151009', sf_config_mock, 1, None, kwargs), ]) mock_task_status.is_completed_task.assert_called_once_with( 'feed_name', '2015-10-09', load_raw_table_tasks.TASK_ID) def test_clear_staging_raw_itunes( sf_config_mock, mock_executor_context, mock_sql_loader, mock_task_status): """Test clear_staging_raw_itunes task.""" load_raw_table_tasks.clear_staging_raw_itunes( activity.Activity(boto3.client('swf', 'us-east-1')), '2015-10-09', 'feed_name', sf_config_mock, 'theorchard', ) mock_executor_context.execute_query.assert_any_call( mock_sql_loader, 'delete_from_staging_raw', mock.ANY) mock_task_status.is_completed_task.assert_called_once_with( 'feed_name', '2015-10-09', load_raw_table_tasks.TASK_ID) def test_load_staging_raw_itunes( mock_config, sf_config_mock, mock_executor_context, sql_loader_version, monkeypatch, mock_task_status): """Test load_staging_raw_itunes task.""" monkeypatch.setattr( garcon_feed_status, 'set_overall_status', MagicMock(return_value=True)) monkeypatch.setattr( itunes_helpers, 'check_status', MagicMock(return_value=True)) monkeypatch.setattr( helpers, 'check_s3_key_exist', MagicMock(return_value=True)) monkeypatch.setattr( s3utils, 'get_file_size', MagicMock(return_value=100)) load_raw_table_tasks.load_staging_raw_itunes( activity.Activity(boto3.client('swf', 'us-east-1')), '2015-10-09', 'itunes_theorchard', sf_config_mock, {'vendor_id': '80029727', 'temp_table': 'staging_raw_itunes_theorchard_80029727_20151009', 'filename': 'D_D_80029727_20151009.txt.gz', 'destination_s3_path': 's3://bucket/iTunes/archives/2015-10-09/' 'D_D_80029727_20151009.txt.gz'} ) assert mock_executor_context.execute_query.call_count == 2 mock_task_status.is_completed_task.assert_called_once_with( 'itunes_theorchard', '2015-10-09', load_raw_table_tasks.TASK_ID) # # Check when status is less than INGESTED # garcon_feed_status.set_overall_status.assert_any_call( # 'iTunes_theorchard', '2015-10-09', 'POPULATED_RAW_TABLE') def test_update_upc_isrc_on_staging_raw( mock_config, mock_sf_config, mock_executor_context, mock_sql_loader, monkeypatch): """Test update_upc_isrc_on_staging_raw.""" load_raw_table_tasks.update_upc_isrc_on_staging_raw( activity.Activity(boto3.client('swf', 'us-east-1')), {}, 'theorchard', '2015-10-09', skip_mapping=None) mock_executor_context.execute_query.assert_any_call( mock_sql_loader, 'update_upc_isrc_staging_raw_theorchard', mock.ANY) def test_update_upc_isrc_on_staging_raw_skip_true( mock_config, mock_sf_config, mock_executor_context, mock_sql_loader, monkeypatch): """Test update_upc_isrc_on_staging_raw.""" load_raw_table_tasks.update_upc_isrc_on_staging_raw( activity.Activity(boto3.client('swf', 'us-east-1')), {}, 'theorchard', '2015-10-09', skip_mapping='True') mock_executor_context.execute_query.assert_not_called()