"""Unit tests for tasks for Amazon Music Ingestion Workflow.""" from datetime import datetime from datetime import timedelta import itertools import os from pathlib import Path import shutil from unittest.mock import call from unittest.mock import MagicMock from unittest.mock import Mock from unittest.mock import mock_open from unittest.mock import patch from garcon_contrib.dynamo_feed_status import garcon_feed_status import pytest from feed_ingestion.flows.amazon_music import config from feed_ingestion.flows.amazon_music import tasks module_path = 'feed_ingestion.flows.amazon_music.tasks' @pytest.mark.parametrize('row, expected', [ # no changes ( '20200625 1 amzn1.account.AFL2SI2XNZQBKTJHHHXD62YRAA5A PCO3 Amazon Prime - Prime Paid Calvin Harris [feat. Big Sean, Katy Perry & Pharrell Williams] Feels GBARL1700772 886446500355 Funk Wav Bounces Vol.1 A10328E0008299868A SS 1 iOS 6 1 20200625T10:06:53 199 34 SONGS 20200625T10:07:27 +10:00 10 11'.split('\t'), # noqa '20200625 1 amzn1.account.AFL2SI2XNZQBKTJHHHXD62YRAA5A PCO3 Amazon Prime - Prime Paid Calvin Harris [feat. Big Sean, Katy Perry & Pharrell Williams] Feels GBARL1700772 886446500355 Funk Wav Bounces Vol.1 A10328E0008299868A SS 1 iOS 6 1 20200625T10:06:53 199 34 SONGS 20200625T10:07:27 +10:00 10 11'.split('\t'), # noqa ), # some sme files have different timestamp format ( '20190202 1 amzn1.account.AG3XT4ZF6BK4G4HC2IYYEXV52CJQ B07N95KF6Y PCO3 Amazon Prime - Prime Paid AUDIO Zoé Veneno MXF140300255 7509951453420 Rocanlover A10328E0002690028O SS 21 1 iOS 6 0 2019-02-02T17:12:32.000Z 1 PRIME_PLAYLIST ParentAsin:B07L528BLH;ChildAsin:B07N95KF6Y 2019-02-02 17:13:08.415 -06:00 10'.split('\t'), # noqa '20190202 1 amzn1.account.AG3XT4ZF6BK4G4HC2IYYEXV52CJQ B07N95KF6Y PCO3 Amazon Prime - Prime Paid AUDIO Zoé Veneno MXF140300255 7509951453420 Rocanlover A10328E0002690028O SS 21 1 iOS 6 0 20190202T17:12:32 1 PRIME_PLAYLIST ParentAsin:B07L528BLH;ChildAsin:B07N95KF6Y 20190202T17:13:08 -06:00 10 \n'.split('\t'), # noqa ), ( '20190202 1 amzn1.account.AECGY53IOGBWXUJV5G5LWVBY7V6Q A2ZU62AUVD5L3T Q174 Amazon Prime - Prime Paid AUDIO Piccolo Coro Mariele Ventre Dell\'Antoniano Il coccodrillo come fa? ITA291700009 886446447384 Il meglio dello zecchino d\'oro (Remastered 2017) A10328E0008267273B SS 8 Fire OS 6 1 2019-02-02T16:16:35.967Z 73 PRIME_STATION SeedType:GENRE;SeedID:2f97f7ac-d96d-433b-ad64-9baae919d36a 2019-02-02 16:16:35.967 10'.split('\t'), # noqa '20190202 1 amzn1.account.AECGY53IOGBWXUJV5G5LWVBY7V6Q A2ZU62AUVD5L3T Q174 Amazon Prime - Prime Paid AUDIO Piccolo Coro Mariele Ventre Dell\'Antoniano Il coccodrillo come fa? ITA291700009 886446447384 Il meglio dello zecchino d\'oro (Remastered 2017) A10328E0008267273B SS 8 Fire OS 6 1 20190202T16:16:35 73 PRIME_STATION SeedType:GENRE;SeedID:2f97f7ac-d96d-433b-ad64-9baae919d36a 20190202T16:16:35 10 \n'.split('\t'), # noqa ), # BUT some times sme files have even more different timestamp format ( '20190202 1 amzn1.account.AECGY53IOGBWXUJV5G5LWVBY7V6Q A2ZU62AUVD5L3T Q174 Amazon Prime - Prime Paid AUDIO Piccolo Coro Mariele Ventre Dell\'Antoniano Il pescecane (solo un ciao) ITA291700080 886446727714 Zecchino D\'Oro 60° Edizione (2017) A10328E0008512864D SS 8 Fire OS 6 1 2019-02-02T20:14:25.000Z 162 PRIME_STATION SeedType:GENRE;SeedID:2f97f7ac-d96d-433b-ad64-9baae919d36a 2019-02-02 20:14:25 10 '.split('\t'), # noqa '20190202 1 amzn1.account.AECGY53IOGBWXUJV5G5LWVBY7V6Q A2ZU62AUVD5L3T Q174 Amazon Prime - Prime Paid AUDIO Piccolo Coro Mariele Ventre Dell\'Antoniano Il pescecane (solo un ciao) ITA291700080 886446727714 Zecchino D\'Oro 60° Edizione (2017) A10328E0008512864D SS 8 Fire OS 6 1 20190202T20:14:25 162 PRIME_STATION SeedType:GENRE;SeedID:2f97f7ac-d96d-433b-ad64-9baae919d36a 20190202T20:14:25 10 '.split('\t'), # noqa ), ( '20171031 1 amzn1.account.AGGNHDBL34QA33MB67Y2CT2HJS7A B075JNF1JS PMO7 Amazon (Subscription) - Free Trials - 30 day trial Rock Toto Africa USSM19801941 884977724950 Africa A10328E0005321444X SS 0 20171030T23:10:55 226 157 UNLIMITED_PLAYLIST ParentAsin:B073PWKXW4;ChildAsin:B075JNF1JS 2017-10-30T23:10:59.000Z +01:00 10 '.split('\t'), # noqa '20171031 1 amzn1.account.AGGNHDBL34QA33MB67Y2CT2HJS7A B075JNF1JS PMO7 Amazon (Subscription) - Free Trials - 30 day trial Rock Toto Africa USSM19801941 884977724950 Africa A10328E0005321444X SS 0 20171030T23:10:55 226 157 UNLIMITED_PLAYLIST ParentAsin:B073PWKXW4;ChildAsin:B075JNF1JS 20171030T23:10:59 +01:00 10 '.split('\t'), # noqa ), ( '20190103 1 amzn1.account.AF5LMTPR5PHX3U3XXDPRFSPS5DKQ Amazon Prime - Prime Paid Soundtracks アレクサンドル・デスプラ A Different Equation USSM11406184 886444761109 The Imitation Game (Original Motion Picture Soundtrack) A10328E0007049348M SS 1 Android 1 2019-01-03 11:20:53.395 175 17 AUTO_PLAYLIST 2019-01-03T11:20:53Z +09:00 10 '.split('\t'), # noqa '20190103 1 amzn1.account.AF5LMTPR5PHX3U3XXDPRFSPS5DKQ Amazon Prime - Prime Paid Soundtracks アレクサンドル・デスプラ A Different Equation USSM11406184 886444761109 The Imitation Game (Original Motion Picture Soundtrack) A10328E0007049348M SS 1 Android 1 20190103T11:20:53 175 17 AUTO_PLAYLIST 20190103T11:20:53 +09:00 10 '.split('\t'), # noqa ), ]) def test__activity_report_row_handler(row, expected): """Test _activity_report_row_handler.""" assert tasks._activity_report_row_handler(row) == expected @pytest.fixture def context(): """Flow context with custom country list.""" return {'activity': MagicMock(), 'date': '2017-08-01', 'report_name': 'test', 'reload': None, 'soft_reload': None} @pytest.fixture def context_reload(): """Flow context with custom country list.""" return {'activity': MagicMock(), 'date': '2017-08-01', 'report_name': 'test', 'reload': 'True', 'soft_reload': None} @pytest.fixture def context_with_custom_countries(): """Flow context with custom country list.""" return {'activity': MagicMock(), 'date': '2017-08-01', 'report_name': 'test', 'reload': None, 'countries': 'US,GB', 'soft_reload': None} @pytest.fixture def context_with_custom_snowflake_error_limit(): """Flow context with custom snowflake_error_limit.""" return {'activity': MagicMock(), 'date': '2017-08-01', 'report_name': 'test', 'reload': None, 'snowflake_error_limit': 10, 'soft_reload': None} # Tests for bootstrap @pytest.fixture def mock_task_status(): """Yield task status.""" task_status_path = ( 'feed_ingestion.flows.amazon_music.tasks.task_status') with patch(task_status_path) as task_status: task_status.is_completed_task.return_value = False yield task_status @pytest.fixture def mock_get_overall_status(): """Yield overall status.""" overall_status_path = ( 'feed_ingestion.flows.amazon_music.tasks.garcon_feed_status.' 'get_overall_status') with patch(overall_status_path) as overall_status: yield overall_status @pytest.fixture def mock_delete_status(): """Yield delete status.""" delete_status_path = ( 'feed_ingestion.flows.amazon_music.tasks.garcon_feed_status.' 'delete_status') with patch(delete_status_path) as delete_status: yield delete_status @pytest.fixture def expected_bootstrap_response(): """Response for bootstrap task.""" return { 'feed_name': 'amazon_music_theorchard_test', 'date': '2017-08-01', 'licensor': 'theorchard', 'report_name': 'test', 'source': 'sftp', 'date_as_in_uuid': '20170801', 'countries': ['AT', 'US'], 'use_s3': False, 'populate_only': False, 'expected_files': sorted([ 'ORCA_US_20170801_Daily_Playlist_Report.txt.zip', 'ORCA_AT_20170801_Daily_Playlist_Report.txt.zip', 'ORED_US_20170801_Daily_Playlist_Report.txt.zip', 'ORED_AT_20170801_Daily_Playlist_Report.txt.zip', 'ORCA_US_20170801_Daily_Activity_Report.txt.zip', 'ORCA_AT_20170801_Daily_Activity_Report.txt.zip', 'ORED_US_20170801_Daily_Activity_Report.txt.zip', 'ORED_AT_20170801_Daily_Activity_Report.txt.zip', 'ORCA_US_20170801_Daily_User_Report.txt.zip', 'ORCA_AT_20170801_Daily_User_Report.txt.zip', 'ORED_US_20170801_Daily_User_Report.txt.zip', 'ORED_AT_20170801_Daily_User_Report.txt.zip']), 'replace_archive_files': True, 's3_drop_location': 's3://dev-feed-drop/feed-drop/AmazonTest/theorchard/', 'archive_bucket': 's3://archive-bucket/AmazonTest/archives/2017-08-01/theorchard/', 'clean_path': 's3://archive-bucket/AmazonTest/archives/2017-08-01' '/theorchard/clean/', 'staging_raw_table': 'staging_raw_amazon_test', 'snowflake_error_limit': 1, 'snowflake_error_on_column_count_mismatch': 'true', 'extract_original_filename': False, 'common_kwargs': { 'licensor': 'theorchard', 'use_s3': False, 'country_list': ['AT', 'US'], 'source': 'sftp'}, 'dimension_tables': config.dimension_tables, 'jenkins_config': config.jenkins_config, 'soft_reload': False } # Tests for bootstrap def test_bootstrap( mock_get_overall_status, mock_delete_status, mock_task_status, expected_bootstrap_response, context, mock_config): """Check that bootstrap returns expected results.""" response = tasks.bootstrap(**context) assert response == expected_bootstrap_response mock_delete_status.assert_not_called() mock_get_overall_status.assert_called_once_with( 'amazon_music_theorchard_test', '2017-08-01') def test_bootstrap_with_reload( mock_config, mock_get_overall_status, mock_delete_status, mock_task_status, expected_bootstrap_response, context_reload): """Check that bootstrap delete overall status.""" response = tasks.bootstrap(**context_reload) assert response == expected_bootstrap_response mock_delete_status.assert_called_once_with( 'amazon_music_theorchard_test', '2017-08-01') mock_get_overall_status.assert_not_called() def test_bootstrap_if_already_ingested( mock_config, mock_get_overall_status, mock_delete_status, mock_task_status, expected_bootstrap_response, context): """Check that bootstrap returns expected results.""" mock_get_overall_status.return_value = garcon_feed_status.STATUS_INGESTED response = tasks.bootstrap(**context) assert response == tasks.STOP_RESPONSE mock_delete_status.assert_not_called() mock_get_overall_status.assert_called_once_with( 'amazon_music_theorchard_test', '2017-08-01') def test_bootstrap_replace_archive_files_when_not_downloaded( mock_config, mock_get_overall_status, mock_delete_status, mock_task_status, expected_bootstrap_response, context): """Archive files are replaced if status is not DOWNLOADED.""" assert tasks.bootstrap(**context).get('replace_archive_files') def test_bootstrap_do_not_replace_archive_files_when_downloaded( mock_config, mock_get_overall_status, mock_delete_status, mock_task_status, expected_bootstrap_response, context): """Archive files are not replaced if status is DOWNLOADED.""" mock_task_status.is_completed_task.return_value = ( 'set_overall_status_DOWNLOADED') assert not tasks.bootstrap(**context).get('replace_archive_files') def test_bootstrap_countries( mock_config, mock_get_overall_status, mock_delete_status, mock_task_status, context): """Check that bootstrap returns expected results.""" assert tasks.bootstrap(**context).get('countries') == ['AT', 'US'] def test_bootstrap_with_custom_countries( mock_config, mock_get_overall_status, mock_delete_status, mock_task_status, context_with_custom_countries): """If country list is not provided then bootstrap uses default list.""" assert tasks.bootstrap( **context_with_custom_countries).get('countries') == ['US', 'GB'] def test_bootstrap_with_custom_snowflake_error_limit( mock_config, context_with_custom_snowflake_error_limit, mock_task_status, mock_get_overall_status): """If country list is not provided then bootstrap uses default list.""" assert tasks.bootstrap( **context_with_custom_snowflake_error_limit).get( 'snowflake_error_limit') == 10 @pytest.mark.parametrize('licensor', ['theorchard', 'sme']) def test_bootstrap_source_datapulse_allowed_licensors( mock_config, mock_get_overall_status, mock_delete_status, mock_task_status, context, licensor): """source=datapulse is allowed for sme and theorchard.""" context['licensor'] = licensor context['source'] = 'datapulse' response = tasks.bootstrap(**context) assert response.get('source') == 'datapulse' assert response.get('licensor') == licensor assert (response.get('staging_raw_table') == config.DATAPULSE_STAGING_RAW_TABLE) @pytest.mark.parametrize('licensor', ['awal', 'altafonte']) def test_bootstrap_source_datapulse_rejects_other_licensors( mock_config, mock_get_overall_status, mock_delete_status, mock_task_status, context, licensor, monkeypatch): """source=datapulse must reject awal/altafonte.""" monkeypatch.setattr( config, 'licensors', ['theorchard', 'sme', 'awal', 'altafonte']) context['licensor'] = licensor context['source'] = 'datapulse' with pytest.raises(AssertionError, match='source="datapulse"'): tasks.bootstrap(**context) def test_bootstrap_unknown_source( mock_config, mock_get_overall_status, mock_delete_status, mock_task_status, context): """An unknown source value is rejected.""" context['source'] = 'bogus' with pytest.raises(AssertionError, match='unsupported source'): tasks.bootstrap(**context) # Tests for load_staging_raw_from_datapulse task @pytest.fixture def datapulse_task_executor_mock(monkeypatch): """Patch registered_executors and helpers used by the datapulse task.""" from feed_ingestion.flows import registered_executors from feed_ingestion.util import task_status executor_with_mock = MagicMock() executor_cm = MagicMock() executor_cm.__enter__.return_value = executor_with_mock executor_class = MagicMock(return_value=executor_cm) monkeypatch.setitem( registered_executors, 'amazon_music_theorchard_unlimited', executor_class) monkeypatch.setattr( task_status, 'is_completed_task', MagicMock(return_value=False)) monkeypatch.setattr( tasks, 'get_sf_config', MagicMock(return_value={'k': 'v'})) return { 'executor_class': executor_class, 'executor': executor_with_mock, } def test_load_staging_raw_from_datapulse_calls_executor( datapulse_task_executor_mock): """Task instantiates the right executor and delegates to it.""" activity = MagicMock() tasks.load_staging_raw_from_datapulse( activity=activity, date='2026-05-22', report_name='unlimited', licensor='theorchard', feed_name='amazon_music_theorchard_unlimited', sfdb_params={'extra': 'param'}, secrets_path='amazon_music', ) datapulse_task_executor_mock['executor'] \ .load_staging_raw_from_datapulse.assert_called_once_with( date='2026-05-22', report_name='unlimited', licensor='theorchard', ) def test_load_staging_raw_from_datapulse_skips_when_already_complete( monkeypatch, datapulse_task_executor_mock): """Task is a no-op when the staging_raw_table_tasks marker is set.""" from feed_ingestion.util import task_status monkeypatch.setattr( task_status, 'is_completed_task', MagicMock(return_value=True)) activity = MagicMock() tasks.load_staging_raw_from_datapulse( activity=activity, date='2026-05-22', report_name='unlimited', licensor='theorchard', feed_name='amazon_music_theorchard_unlimited', sfdb_params={}, secrets_path='amazon_music', ) datapulse_task_executor_mock['executor_class'].assert_not_called() datapulse_task_executor_mock['executor'] \ .load_staging_raw_from_datapulse.assert_not_called() class TestDownloadAndClean(object): """Test download_and_clean.""" date = '2017-08-01' file_key = '2017-08-01/so-say-we-all.txt.zip' file_key_path = 'file_stage_test_theorchard/2017-08-01' \ '/so-say-we-all.txt.zip' local_file = 'file_stage_test_theorchard/2017-08-01/so-say-we-all.txt.zip' unzipped_file = 'file_stage_test_theorchard/2017-08-01/so-say-we-all.txt' new_file = 'file_stage_test_theorchard/2017-08-01/clean/so-say-we-all.txt' gzipped_file = 'file_stage_test_theorchard/2017-08-01/so-say-we-all.txt.gz' archive_bucket = 'some-bucket' report = 'test' @pytest.fixture def mock_resource_call(self): """Return a mock.""" return MagicMock() @pytest.fixture def mock_bucket(self): """Return a mock.""" return MagicMock() @pytest.fixture def mock_zipfile_instance(self): """Return a mock.""" return MagicMock() @pytest.fixture def mock_zipfile(self, mock_zipfile_instance): """Yield the zipfile mock.""" zipfile_path = module_path + '.zipfile' with patch(zipfile_path) as zipfile: zipfile.ZipFile.return_value = mock_zipfile_instance yield zipfile @pytest.fixture def mock_gzip(self): """Yield the zipfile mock.""" gzip_path = module_path + '.gzip' with patch(gzip_path) as gzip: yield gzip @pytest.fixture def mock_s3_resource(self, mock_resource_call, mock_bucket): """Yield the resource mock.""" boto3_resource_path = module_path + '.boto3.resource' with patch(boto3_resource_path) as boto3_resource: mock_resource_call.Bucket.return_value = mock_bucket boto3_resource.return_value = mock_resource_call yield boto3_resource @pytest.fixture def mock_open_file(self): """Yield the open file mock.""" mock_open_file = mock_open() with patch('builtins.open', mock_open_file): yield mock_open_file @pytest.fixture def mock_makedirs(self): """Yield the mkdir mock.""" makedirs_path = module_path + '.os.makedirs' with patch(makedirs_path) as mkdir: yield mkdir @pytest.fixture def mock_mkdir(self): """Yield the mkdir mock.""" mkdir_path = module_path + '.os.mkdir' with patch(mkdir_path) as mkdir: yield mkdir @pytest.fixture def mock_activity(self): """Return mock activity.""" return MagicMock() @pytest.fixture def mock_os_remove(self): """Yield the os mock.""" remove_path = module_path + '.os.remove' with patch(remove_path) as remove: yield remove @pytest.fixture def run_download_and_clean( self, mock_activity, mock_s3_resource, mock_zipfile, mock_open_file, mock_makedirs, mock_mkdir, mock_gzip, mock_os_remove): """Run and return the download_and_clean task.""" return tasks.download_and_clean( mock_activity, self.date, self.file_key, self.archive_bucket, self.report) def test_creates_s3_resource( self, run_download_and_clean, mock_s3_resource): """s3 resource is created.""" mock_s3_resource.assert_called_with('s3') def test_creates_bucket( self, run_download_and_clean, mock_resource_call): """Create the bucket object for the arhive bucket.""" mock_resource_call.Bucket.assert_called_with( self.archive_bucket) def test_downloads_file_to_destination( self, run_download_and_clean, mock_bucket): """Download files locally.""" mock_bucket.download_file.assert_called_with( self.file_key, self.local_file) def test_creates_zipfile(self, run_download_and_clean, mock_zipfile): """Create zipfile object.""" mock_zipfile.ZipFile.assert_called_with(self.local_file) def test_unzips_file(self, run_download_and_clean, mock_zipfile_instance): """Unzip file.""" mock_zipfile_instance.__enter__.return_value.extractall.assert_called() def test_opens_original_file_for_reading( self, run_download_and_clean, mock_open_file): """Open the original file to be read.""" mock_open_file.assert_any_call(self.unzipped_file) def test_makes_staging_dir(self, run_download_and_clean, mock_makedirs): """Create the staging directory for file work.""" mock_makedirs.assert_any_call('file_stage_test_theorchard/2017-08-01', exist_ok=True) def test_creates_gzip_file(self, run_download_and_clean, mock_gzip): """Create gzip file.""" mock_gzip.open.assert_called_with( 'file_stage_test_theorchard/2017-08-01/so-say-we-all.txt.gz', 'wb') def test_uploads_file_to_s3( self, run_download_and_clean, mock_bucket): """Upload new zipped file to s3.""" mock_bucket.upload_file.assert_called_with( 'file_stage_test_theorchard/2017-08-01/so-say-we-all.txt.gz', '2017-08-01/clean/so-say-we-all.txt.gz') def test_logs_some_info_when_downloaded_and_cleaned( self, run_download_and_clean, mock_activity): """Log some info when tables are created and dropped.""" assert mock_activity.logger.info.call_count == 1 def test__prepare_playlist_report_line_v1(self): """Test line preparation for initial playlist report version.""" original_line = 'v1#*#v2#*#v3#*#v4#*#v5\n' expected_value = 'v1\tv2\tv3\tv4\tv5\t\t\n' result = tasks._prepare_report_line('playlist', original_line) assert result == expected_value def test__prepare_playlist_report_line_v2(self): """Test line preparation for v2 playlist report version.""" original_line = 'v1#*#v2#*#v3#*#v4#*#v5#*#v6#*#v7\n' expected_value = 'v1\tv2\tv3\tv4\tv5\tv6\tv7\n' result = tasks._prepare_report_line('playlist', original_line) assert result == expected_value def test__prepare_activity_report_line_non_standard(self): """Test line preparation for initial activiy report.""" non_standard_row = ['' for _ in range(28)] non_standard_timestamp = '20181212T03:09:17' non_standard_row[20] = non_standard_timestamp non_standard_row[25] = non_standard_timestamp standard_row = ['' for _ in range(29)] standard_row[20] = non_standard_timestamp standard_row[25] = non_standard_timestamp original_line = '#*#'.join(non_standard_row) + '\n' expected_value = '\t'.join(standard_row) + '\n' result = tasks._prepare_report_line('activity', original_line) assert result == expected_value def test_remove_local_files( self, run_download_and_clean, mock_os_remove): """Upload new zipped file to s3.""" expected = [ call(self.file_key_path), call(self.unzipped_file), call(self.gzipped_file)] assert mock_os_remove.call_args_list == expected class S3ClientMock(object): """Fake s3 client for SME.""" def list_objects(self, **kwargs): """Return list of bucket contents.""" return { 'Contents': [ { 'Key': f'{TestMapExpectedFilesToSme.sme_prefix}/{file}' } for file in TestMapExpectedFilesToSme.sme_files_in_directory] } class TestMapExpectedFilesToSme(object): """Test map_expected_files_to_sme task.""" the_date = datetime(year=2020, month=4, day=15) date = the_date.strftime('%Y-%m-%d') report_type = 'prime' expected_files = [ 'PCO3_AU_20200415_Daily_Activity_Report.txt.zip', 'PCO3_AU_20200415_Daily_Playlist_Report.txt.zip', 'PCO3_AU_20200415_Daily_User_Report.txt.zip' ] source_bucket_name = 'sme_test_bucket' sme_prefix = f'amazon/in/{report_type}/{the_date.strftime("%Y/%m/%d")}' s3_drop_location = 's3://{bucket}/{prefix}'.format( bucket=source_bucket_name, prefix=sme_prefix) sme_files_in_directory = [ 'PCO3_A_AU_20200415_20200415_Activity_20200416.txt.zip', 'PCO3_A_AU_20200415_20200415_Playlist_20200416.txt.zip', 'PCO3_A_AU_20200415_20200415_User_20200416.txt.zip', ] def run_task(self, mock_activity): """Run and return the download_and_clean task.""" return tasks.map_expected_files_to_sme( activity=mock_activity, date=self.date, s3_drop_location=self.s3_drop_location, expected_files=self.expected_files, ) @pytest.fixture def mock_activity(self): """Return mock activity.""" return MagicMock() def test_mapping(self, mock_activity): """Test of map_expected_files_to_sme task.""" with patch(module_path + '._get_sme_s3_client') as sme_s3_client_mock: sme_s3_client_mock.return_value = S3ClientMock() result = self.run_task(mock_activity) expected_result = { 'map': { 'PCO3_AU_20200415_Daily_Activity_Report.txt.zip': 'amazon/in/prime/2020/04/15/PCO3_A_AU_20200415_20200415_Activity_20200416.txt.zip', # noqa: E501 'PCO3_AU_20200415_Daily_Playlist_Report.txt.zip': 'amazon/in/prime/2020/04/15/PCO3_A_AU_20200415_20200415_Playlist_20200416.txt.zip', # noqa: E501 'PCO3_AU_20200415_Daily_User_Report.txt.zip': 'amazon/in/prime/2020/04/15/PCO3_A_AU_20200415_20200415_User_20200416.txt.zip' # noqa: E501 }, 'has_files': True, } assert result == expected_result class TestSmeGrabAndClean(object): """Test sme_grab_and_clean task.""" the_date = datetime(year=2020, month=4, day=15) date = the_date.strftime('%Y-%m-%d') report_type = 'prime' source_key_name = f'amazon/in/{report_type}' \ f'/{the_date.strftime("%Y/%m/%d")}' \ f'/PCO3_A_AU_20200415_20200415_Activity_20200416.txt.zip' destination_key_name = f'AmazonPrime/archives/{date}' \ f'/sme' \ f'/PCO3_AU_20200415_Daily_Activity_Report.txt.zip' destination_bucket_name = 'test-cucumbers' sme_bucket = 'test_sme_bucket' feed_name = 'amazon_music' @pytest.fixture def mock_zipfile_instance(self): """Return a mock.""" return MagicMock() @pytest.fixture def mock_zipfile(self, mock_zipfile_instance): """Yield the zipfile mock.""" zipfile_path = module_path + '.zipfile' with patch(zipfile_path) as zipfile: zipfile.ZipFile.return_value = mock_zipfile_instance yield zipfile @pytest.fixture def mock_gzip(self): """Yield the zipfile mock.""" gzip_path = module_path + '.gzip' with patch(gzip_path) as gzip: yield gzip @pytest.fixture def mock_resource_call(self): """Return a mock.""" return MagicMock() @pytest.fixture def mock_s3_resource(self, mock_resource_call, mock_bucket): """Yield the resource mock.""" boto3_resource_path = module_path + '.boto3.resource' with patch(boto3_resource_path) as boto3_resource: mock_resource_call.Bucket.return_value = mock_bucket boto3_resource.return_value = mock_resource_call yield boto3_resource @pytest.fixture def mock_open_file(self): """Yield the open file mock.""" mock_open_file = mock_open() with patch('builtins.open', mock_open_file): yield mock_open_file @pytest.fixture def mock_bucket(self): """Return a mock.""" return MagicMock() @pytest.fixture def mock_garcon_feed_status(self): """Mock garcon_feed_status.""" garcon_feed_status_path = module_path + '.garcon_feed_status' with patch(garcon_feed_status_path) as garcon_feed_status: garcon_feed_status.get_status.return_value = None yield garcon_feed_status_path @pytest.fixture def mock_garcon_feed_status_downloaded(self): """Mock garcon_feed_status.""" garcon_feed_status_path = module_path + '.garcon_feed_status' with patch(garcon_feed_status_path) as garcon_feed_status: garcon_feed_status.get_status.return_value \ = garcon_feed_status.STATUS_DOWNLOADED yield garcon_feed_status_path @pytest.fixture def mock_makedirs(self): """Yield the mkdir mock.""" makedirs_path = module_path + '.os.makedirs' with patch(makedirs_path) as mkdir: yield mkdir @pytest.fixture def mock_mkdir(self): """Yield the mkdir mock.""" mkdir_path = module_path + '.os.mkdir' with patch(mkdir_path) as mkdir: yield mkdir @pytest.fixture def mock_activity(self): """Return mock activity.""" return MagicMock() @pytest.fixture def mock_os_remove(self): """Yield the os mock.""" remove_path = module_path + '.os.remove' with patch(remove_path) as remove: yield remove @pytest.fixture def mock_sme_s3_client(self): """Yield sme s3 client mock.""" with patch(module_path + '._get_sme_s3_client') as mock: yield mock @pytest.fixture def run_sme_grab_and_clean( self, mock_activity, mock_s3_resource, mock_zipfile, mock_open_file, mock_makedirs, mock_mkdir, mock_gzip, mock_os_remove, mock_sme_s3_client ): """Run and return the download_and_clean task.""" return tasks.sme_grab_and_clean( activity=mock_activity, report_name=self.report_type, date=self.date, feed_name=self.feed_name, source_bucket_name=self.sme_bucket, source_key_name=self.source_key_name, destination_bucket_name=self.destination_bucket_name, destination_key_name=self.destination_key_name, ) def test_success_return_value( self, mock_garcon_feed_status, mock_s3_resource, run_sme_grab_and_clean ): """Test return value for success execution.""" assert run_sme_grab_and_clean == {'success': True} def test_skip( self, mock_garcon_feed_status_downloaded, run_sme_grab_and_clean, ): """Test return value for skip execution.""" assert run_sme_grab_and_clean == {'skip': True} class TestAwalAltafonteGrabDropFiles(object): """Test awal_altafonte_grab_drop_files task.""" the_date = datetime(year=2020, month=4, day=15) date = the_date.strftime('%Y-%m-%d') report_name = 'unlimited' destination_key_name = ( 'AmazonUnlimitedV2/archives/2020-04-15/awal/' 'ZQAWA_Daily_Unlimited_Activity_20200415_AU.txt.zip') s3_archive_path = ( 's3://dev-cucumbers/AmazonUnlimitedV2/archives/2020-04-15/awal/') filename = 'ZQAWA_Daily_Unlimited_Activity_20200415_AU.txt.zip' @pytest.fixture def mock_zipfile_instance(self): """Return a mock.""" return MagicMock() @pytest.fixture def mock_garcon_ftp(self): """Yield the garcon_ftp mock.""" garcon_ftp_path = module_path + '.garcon_ftp' with patch(garcon_ftp_path) as garconftp: copy_file_from_ftp_to_s3_mock = MagicMock() garconftp.copy_file_from_ftp_to_s3 = copy_file_from_ftp_to_s3_mock yield copy_file_from_ftp_to_s3_mock @pytest.fixture def mock_get_secret(self): """Yield get_secret.""" path = module_path + '.get_secret' with patch(path) as mock_secret: mock_secret.return_value = 'secret' yield mock_secret def test_return_success_response(self, mock_garcon_ftp, mock_get_secret): """Test return value for success execution.""" ftp_response = { 'file': self.filename, 'status': True, 'file_size': 1} mock_garcon_ftp.return_value = ftp_response response = tasks.awal_grab_drop_files( activity=MagicMock(), report_name=self.report_name, date=self.date, destination_key_name=self.destination_key_name, s3_archive_path=self.s3_archive_path ) assert response == ftp_response def test_when_there_is_no_files(self, mock_garcon_ftp, mock_get_secret): """Test return value if there is no such file.""" ftp_response = { 'file': self.filename, 'status': False, 'exception': FileNotFoundError(f'{self.filename}')} mock_garcon_ftp.return_value = ftp_response response = tasks.awal_grab_drop_files( activity=MagicMock(), report_name=self.report_name, date=self.date, destination_key_name=self.destination_key_name, s3_archive_path=self.s3_archive_path ) assert response == ftp_response class TestRemoveFileStage(object): """Test remove_file_stage task.""" report_date = '2017-08-02' report = 'test' @pytest.fixture def mock_rmtree(self): """Yield the rmtree mock.""" rmtree_path = module_path + '.shutil.rmtree' with patch(rmtree_path) as rmtree: yield rmtree @pytest.fixture def mock_os_path_exists(self): """Yield the os.path.exists mock.""" with patch('os.path.exists') as mock: yield mock @pytest.fixture def run_remove_file_stage(self, mock_rmtree): """Run and return the remove_file_stage task.""" return tasks.remove_file_stage( MagicMock(), self.report_date, self.report, 'theorchard') def test_removes_file_stage_dir(self, mock_os_path_exists, run_remove_file_stage, mock_rmtree): """Remove file stage.""" mock_rmtree.assert_called_with('file_stage_test_theorchard/2017-08-02') class TestCheckForNewFiles(object): """Test check_for_new_files.""" @pytest.fixture def common_files(self): """List of regular files.""" return [ 'ORCA_US_20170801_Daily_Playlist_Report.txt.zip', 'ORCA_US_20170801_Daily_User_Report.txt.zip', 'ORCA_AT_20170801_Daily_Playlist_Report.txt.zip', 'ORCA_AT_20170801_Daily_User_Report.txt.zip'] @pytest.fixture def unexpected_files(self): """List of files extended with new countries.""" return [ 'ORCA_MX_20170801_Daily_Playlist_Report.txt.zip', 'unexpected_file.txt'] @pytest.fixture def unstable_countries_files(self): """List of files extended with new countries.""" return [ 'ORCA_CA_20170801_Daily_Playlist_Report.txt.zip', 'ORCA_ROE_EU_20170801_Daily_Playlist_Report.txt.zip'] @pytest.fixture def mock_unstable_countries(self, mocker): """Mock of stability monitoring countries.""" countries = ['ES', 'US', 'CA', 'ROE_EU'] mocker.patch( 'feed_ingestion.flows.amazon_music.config.reports', { 'test': { 'unstable_countries': countries, # todo: remove as nused 'licensor_org_country': { 'theorchard': { 'ORCA': countries, 'ORED': countries, }, }, }}) return countries @pytest.fixture def stable_countries(self, mock_unstable_countries): """List of stable countries for tests.""" return mock_unstable_countries[:2] @pytest.fixture def unstable_countries(self, mock_unstable_countries): """List of unstable countries for tests.""" return mock_unstable_countries[2:] @pytest.fixture def stable_country_files(self, stable_countries): """Full list of required files for country stability.""" file_template = ( '{org}_{country}_{date}_Daily_{report_type}_Report.txt.zip') file_params = list(itertools.product( stable_countries, config.report_types)) files = [] today = datetime.today() for date_shift in range( config.MONITORING_PERIOD_OF_UNSTABLE_COUNTRIES): date = today - timedelta(days=config.FEED_DELAY + date_shift) for country, report_type in file_params: for org in ['ORCA', 'ORED']: files.append(file_template.format( org=org, country=country, report_type=report_type, date=date.strftime('%Y%m%d'))) return files @patch('feed_ingestion.util.aws.s3.get_list_of_files_and_directories') def test_filter_files( self, list_s3, mock_config, common_files, unexpected_files, unstable_countries_files): """Test _filter_files function.""" list_s3.return_value = ( common_files + unexpected_files + unstable_countries_files) date = '2020-01-30' res_unexpected_files, res_stability_monitoring_files = ( tasks._filter_files(date, 's3_path', 'theorchard', 'test')) assert res_unexpected_files == unexpected_files assert res_stability_monitoring_files == unstable_countries_files @patch('feed_ingestion.flows.amazon_music.tasks._filter_files') def test_monitor_drop_location_unexpected(self, filter_files_mock, mock_config): """Test monitor_drop_location with unexpected files.""" filter_files_mock.return_value = (['unexpected1', 'unexpected2'], None) date = '2020-01-30' result = tasks.monitor_drop_location(Mock(), date, 's3_location', 'report', 'theorchard') assert 'unexpected_files_subject' in result @patch('feed_ingestion.flows.amazon_music.tasks._filter_files') def test_monitor_drop_location_common(self, filter_files_mock): """Test monitor_drop_location with regular files only.""" filter_files_mock.return_value = None, None date = '2020-01-30' result = tasks.monitor_drop_location(Mock(), date, 's3_location', 'report', 'theorchard') assert result == {} @patch('feed_ingestion.flows.amazon_music.tasks._filter_files') def test_monitor_drop_location_stability( self, filter_files_mock, stable_country_files, mock_countries_config): """Test monitor_drop_location with unexpected files.""" filter_files_mock.return_value = (None, stable_country_files,) date = '2020-01-30' result = tasks.monitor_drop_location(Mock(), date, 's3_location', 'test', 'theorchard') assert 'new_stable_countries_subject' in result @patch('feed_ingestion.flows.amazon_music.tasks._filter_files') def test_monitor_drop_location_skip_monitor( self, filter_files_mock, mock_config, stable_country_files, mock_unstable_countries): """Test monitor_drop_location with unexpected files.""" filter_files_mock.return_value = (None, stable_country_files,) date = '2020-01-30' result = tasks.monitor_drop_location( activity=Mock(), date=date, s3_drop_location='s3_location', report_name='test', licensor='theorchard', skip_monitor='True') assert result == {'skip_monitor': 'True'} def test_evaluate_stability(self, stable_country_files, stable_countries, unstable_countries, mock_countries_config): """Test _evaluate_stability function logic.""" result = tasks._evaluate_stability(stable_country_files, 'theorchard', 'test') for country in stable_countries: assert result[country] for country in unstable_countries: assert not result[country] def test_new_stable_countries_message(self): """Test _stability_message function.""" countries_stability = {'ES': True, 'CA': False} country_statuses = [ '{} {}'.format(country, 'stable' if stable else 'unstable') for country, stable in countries_stability.items() ] expected_result = ( config.NEW_STABLE_COUNTRIES_NOTIFICATION_MESSAGE.format('\n'.join( country_statuses))) result = tasks._new_stable_countries_message(countries_stability) assert result == expected_result @patch('feed_ingestion.flows.amazon_music.tasks.' '_new_stable_countries_message') @patch('feed_ingestion.flows.amazon_music.tasks._evaluate_stability') def test_unstable_countries_report( self, evaluate_stability, new_stable_countries_message): """Test _stability_files_report when the report is not empty.""" evaluate_stability.return_value = {'ES': True, 'CA': False} new_stable_countries_message.return_value = 'Test message' result = tasks._unstable_countries_report(MagicMock(), 'theorchard', 'report') assert 'new_stable_countries_subject' in result assert 'new_stable_countries_message' in result @patch('feed_ingestion.flows.amazon_music.tasks.' '_new_stable_countries_message') @patch('feed_ingestion.flows.amazon_music.tasks._evaluate_stability') def test_unstable_countries_report_empty( self, evaluate_stability, new_stable_countries_message): """Test _stability_files_report when there is no data for report.""" evaluate_stability.return_value = {'ES': False, 'CA': False} result = tasks._unstable_countries_report(MagicMock(), 'theorchard', 'report') assert result == {} class TestHelperFunctions: """Test helper functions the in task module.""" @pytest.mark.parametrize( 'countries, expected_result', [ ('FR', ['FR']), ('ES,IT', ['ES', 'IT']), (' US, GB ', ['US', 'GB']) ]) def test_get_context_countries_with_valid_value( self, countries, expected_result): """Test _get_context_countries with valid values.""" assert tasks._get_context_countries(countries) == expected_result @pytest.mark.parametrize( 'countries', [ 'es,it', 'ES IT', 'USA, IT', ]) def test_get_context_countries_with_invalid_invalue(self, countries): """Test _get_context_countries with invalid values.""" with pytest.raises(ValueError): tasks._get_context_countries(countries) @pytest.fixture def mock_executor_context(): """Yield executor context.""" reg_executors_class_path = ( 'feed_ingestion.flows.amazon_music.tasks.registered_executors') with patch(reg_executors_class_path) as reg_executors: mock_executor = MagicMock() reg_executors.get = mock_executor yield mock_executor.return_value.return_value.__enter__.return_value @pytest.fixture def mock_check_status(): """Yield task status.""" task_status_path = 'feed_ingestion.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 @pytest.fixture def stage_files_fixture(mock_config): """Create and yield tuple of 2 files.""" staging_dir = tasks.STAGING_PATH.format( licensor='theorchard', report='test_report', ) def create_file(file, mtime): dirname = os.path.dirname(file) os.makedirs(dirname, exist_ok=True) Path(file).touch(exist_ok=True) os.utime(file, (mtime, mtime)) staled_file = os.path.join(staging_dir, '2019-01-01', 'staled.txt') fresh_file = os.path.join(staging_dir, '2019-01-02', 'fresh.txt') now_epoch = datetime.now().timestamp() # in mock_config MAX_EXECUTION_TIMEOUT = 3600 create_file( file=staled_file, mtime=now_epoch - 3600 * 2) create_file( file=fresh_file, mtime=now_epoch - 3000) yield staled_file, fresh_file shutil.rmtree(staging_dir) def test_remove_stale_stage_files(stage_files_fixture): """Test remove_stale_stage_files.""" staled_file, fresh_file = stage_files_fixture assert os.path.exists(staled_file) assert os.path.exists(fresh_file) tasks.remove_stale_stage_files( activity=MagicMock(), report_name='test_report', licensor='theorchard', ) assert not os.path.exists(staled_file) assert os.path.exists(fresh_file) def test_load_aggregated_table_table( mock_check_status, mock_executor_context, mock_config): """Test test_load_demographics_table.""" date = '2017-08-01' sfdb_params = {'db': 'db', 'schema': 'schema'} tasks.load_aggregated_table( MagicMock(), date, 'feed_name', 'theorchard', 'test', sfdb_params) mock_executor_context.clean_aggregated_staging_raw_table. \ assert_called_with(date, ['ORCA', 'ORED'], source=None) mock_executor_context.load_aggregated_staging_raw_table.assert_called_with( date, source=None) def test_load_aggregated_table_datapulse( mock_check_status, mock_executor_context, mock_config): """Datapulse skips store-code orgs and scopes by feedid + licensor.""" date = '2017-08-01' sfdb_params = {'db': 'db', 'schema': 'schema'} tasks.load_aggregated_table( MagicMock(), date, 'feed_name', 'theorchard', 'test', sfdb_params, source=config.SOURCE_DATAPULSE) mock_executor_context.clean_aggregated_staging_raw_table. \ assert_called_with(date, None, source=config.SOURCE_DATAPULSE) mock_executor_context.load_aggregated_staging_raw_table.assert_called_with( date, source=config.SOURCE_DATAPULSE)