"""Unit tests for tasks for Amazon Digital Services Ingestion Workflow.""" import datetime import os from unittest import mock from unittest.mock import call from unittest.mock import MagicMock from unittest.mock import patch from freezegun import freeze_time import pytest from feed_ingestion.flows.amazon_digital_services import config from feed_ingestion.flows.amazon_digital_services import tasks from tests.conftest import SubstringMatcher @pytest.fixture def mock_config_ads(monkeypatch): """Mock config Amazon Digital Services.""" config_values = { 'vendors_countries': { 'theorchard': [ (datetime.date(2017, 1, 1), { 'OR': { 'required': 'AT,US'.split(','), 'optional': 'CH,DE,JP'.split(','), }, 'RED': { 'required': [], 'optional': 'AT,CA,JP,US'.split(','), }, }), ], 'sme': [ (datetime.date(2016, 12, 4), { 'PLV1': { 'required': ['US'], 'optional': [], }, }), (datetime.date(2018, 1, 1), { 'PLV1': { 'required': ['US'], 'optional': [], }, 'PJ20': { 'required': ['AT'], 'optional': ['IT'], }, 'P748': { 'required': ['US'], 'optional': [], }, }), ], }, } for attr_name, value in config_values.items(): monkeypatch.setattr(config, attr_name, value) yield config @pytest.fixture def mock_boto3(): """Mock boto3.""" boto3_path = ( 'feed_ingestion.flows.amazon_digital_services.tasks.boto3') with patch(boto3_path) as boto3: mock_client = MagicMock() mock_client.send_email = MagicMock() boto3.client.return_value = mock_client yield boto3 @pytest.fixture() def ftp_conf_mock(monkeypatch): """Mock ftp configuration.""" sftp = { 'theorchard': { 'host': 'some_host', 'username': 'some_username', 'password': 'some_password', 'port': 22, 'path': ''} } monkeypatch.setattr(config, 'sftp', sftp) return sftp @patch('feed_ingestion.flows.amazon_digital_services.tasks.zipfile') def test__unzip_file(zipfile_mock): """Test _unzip_file util function.""" zip_obj = MagicMock() zip_obj.orig_filename = 'superfile.csv' zipfile_mock.ZipFile.return_value.__enter__.return_value.infolist. \ return_value = [zip_obj] uncompressed_path = tasks._unzip_file('path/file.zip') assert uncompressed_path == 'path/superfile.csv' zipfile_mock.ZipFile.assert_called_with('path/file.zip') @patch('feed_ingestion.flows.amazon_digital_services.tasks.subprocess') def test__remove_footer(subprocess_mock): """Test _unzip_file util function.""" tasks._remove_footer('path/file.csv', lines_to_remove=1) subprocess_mock.check_call.assert_called_with( SubstringMatcher(containing=['sed -e']), shell=True) @pytest.mark.parametrize('context,expected', [ ( { 'activity': MagicMock(), 'date': None, 'licensor': 'theorchard', 'reload': False, }, { 'date': '2000-01-01', 'feed_name': 'amazon_digital_services_theorchard', 'secrets_path': config.secrets_path, 's3_archive_path': 's3://dev-cucumbers/AmazonV2' '/archives/2000-01-01/theorchard/', 's3_tmp_path': 's3://dev-cucumbers/AmazonV2/tmp' '/2000-01-01/theorchard/', 'file_pattern': r'.+(ZQ.+_20000101_Summary_Statement_\w+\.csv\.zip)|' r'.+(ZQ.+_AmazonMP3_JP_Daily_20000101\.txt\.zip)', 'licensor': 'theorchard', 'staging_raw_table': 'staging_raw_amazon_v2' } ), ( { 'activity': MagicMock(), 'date': '2001-01-01', 'licensor': 'theorchard', 'reload': False, }, { 'date': '2001-01-01', 'feed_name': 'amazon_digital_services_theorchard', 'secrets_path': config.secrets_path, 's3_archive_path': 's3://dev-cucumbers/AmazonV2' '/archives/2001-01-01/theorchard/', 's3_tmp_path': 's3://dev-cucumbers/AmazonV2/tmp' '/2001-01-01/theorchard/', 'file_pattern': r'.+(ZQ.+_20010101_Summary_Statement_\w+\.csv\.zip)|' r'.+(ZQ.+_AmazonMP3_JP_Daily_20010101\.txt\.zip)', 'licensor': 'theorchard', 'staging_raw_table': 'staging_raw_amazon_v2' } ), ( { 'activity': MagicMock(), 'date': '2001-01-01', 'reload': False, 'licensor': 'sme' }, { 'date': '2001-01-01', 'feed_name': 'amazon_digital_services_sme', 'secrets_path': config.secrets_path, 's3_archive_path': 's3://dev-cucumbers/AmazonV2/archives/2001-01-01/sme/', 's3_tmp_path': 's3://dev-cucumbers/AmazonV2' '/tmp/2001-01-01/sme/', 'file_pattern': '.+(.+_D_20010101_20010101_.+\\.txt)', 'licensor': 'sme', 'staging_raw_table': 'staging_raw_amazon_v2_sme' } ), ]) @freeze_time('2000-01-01') @patch('feed_ingestion.flows.amazon_digital_services.tasks.garcon_feed_status') def test_bootstrap(feed_status_mock, monkeypatch, context, expected, mock_config_ads): """Test bootstrap.""" reply = tasks.bootstrap(**context) assert reply == expected @pytest.mark.parametrize('licensor, vendor, country, date, is_required', [ ('theorchard', 'OR', 'US', datetime.date(2020, 1, 1), True), ('theorchard', 'RED', 'US', datetime.date(2020, 1, 1), False), ('theorchard', 'RED', 'US', datetime.date(2016, 12, 31), ValueError), ('sme', 'PLV1', 'US', datetime.date(2020, 1, 1), True), ('sme', 'PLV1', 'US', datetime.date(2017, 12, 31), True), ('sme', 'PJ20', 'AT', datetime.date(2020, 1, 1), True), ('sme', 'PJ20', 'AT', datetime.date(2017, 12, 31), False), ]) def test__is_required_vendor_country_for_date( licensor, vendor, country, date, is_required, mock_config_ads): """Test _is_required_vendor_country_for_date().""" params = dict( vendor=vendor, country=country, licensor=licensor, date_obj=date, ) if is_required == ValueError: with pytest.raises(is_required): tasks._is_required_vendor_country_for_date(**params) else: assert tasks._is_required_vendor_country_for_date(**params) \ == is_required @patch( 'feed_ingestion.flows.amazon_digital_services.tasks.garcon_ftp') def test_fetch_from_drop_location_theorchard(ftp_mock, ftp_conf_mock, mock_config_ads): """Test custom fetch_from_drop_location task.""" activity_mock = MagicMock() tasks.fetch_from_drop_location( activity_mock, '2018-02-20', 'amazon_digital_services', 's3://somepath/tmp/', 'theorchard', use_s3=None) copy_file_from_ftp_to_s3_mock = MagicMock() copy_file_from_ftp_to_s3_mock.return_value = { 'file': 'awesome_file', 'status': True, 'file_size': 1000 } ftp_mock.copy_file_from_ftp_to_s3 = copy_file_from_ftp_to_s3_mock activity_mock = MagicMock() result = tasks.fetch_from_drop_location( activity_mock, '2018-02-20', 'amazon_digital_services', 's3://somepath/tmp/', 'theorchard', use_s3=None) assert len(result['source_files_dict']['files']) == 9 assert all(file['found'] for file in result['source_files_dict']['files']) assert copy_file_from_ftp_to_s3_mock.call_count == 9 # check two calls for example copy_file_from_ftp_to_s3_mock.assert_has_calls([ call( activity=activity_mock, ftp_creds=ftp_conf_mock['theorchard'], ftp_path='sales/AT/Daily/', ftp_file_name='ZQOR_20180220_Summary_Statement.csv.zip', s3_path='s3://somepath/tmp/', s3_file_name='ZQOR_20180220_Summary_Statement_AT.csv.zip' ), call( activity=activity_mock, ftp_creds=ftp_conf_mock['theorchard'], ftp_path='sales/JP/Daily/', ftp_file_name='ZQRED_AmazonMP3_JP_Daily_20180220.txt.zip', s3_path='s3://somepath/tmp/', s3_file_name='ZQRED_AmazonMP3_JP_Daily_20180220.txt.zip' ), ], any_order=True) @patch('feed_ingestion.flows.amazon_digital_services.tasks.s3_tasks') @patch('feed_ingestion.flows.amazon_digital_services.tasks.os') def test_fetch_from_drop_location_sme(os_mock, s3_tasks_mock, mock_config_ads): """Test custom fetch_from_drop_location task.""" def download_side_effect(destination_key_name, **kwargs): file_name = destination_key_name.split('/')[-1] return { file_name: True, 'file_size': 100, } download_mock = MagicMock() download_mock.side_effect = download_side_effect s3_tasks_mock.copy_file_from_sme_s3_to_theocrhard = download_mock os_mock.path.getsize.return_value = 10 activity_mock = MagicMock() result = tasks.fetch_from_drop_location( activity_mock, '2018-02-20', 'amazon_digital_services', 's3://somepath/tmp/', 'sme', use_s3=None) assert len(result['source_files_dict']['files']) == 4 assert all(file['found'] for file in result['source_files_dict']['files']) assert download_mock.call_count == 4 # check two calls for example download_mock.assert_has_calls([ call( activity=activity_mock, secrets_path='swf-sme-ca-prod-partners-s3-access', source_bucket_name='sme-ca-prod-partners', source_key_name='amazon/in/download/P748/P748_D_20180220_20180220_US.txt', # noqa:E501 destination_bucket_name='dev-cucumbers', destination_key_name='tmp/P748_D_20180220_20180220_US.txt', # noqa:E501 replace=True, ), call( activity=activity_mock, secrets_path='swf-sme-ca-prod-partners-s3-access', source_bucket_name='sme-ca-prod-partners', source_key_name='amazon/in/download/PJ20/PJ20_D_20180220_20180220_IT.txt', # noqa:E501 destination_bucket_name='dev-cucumbers', destination_key_name='tmp/PJ20_D_20180220_20180220_IT.txt', # noqa:E501 replace=True, ), ], any_order=True) @patch( 'feed_ingestion.flows.amazon_digital_services.tasks.garcon_ftp') @patch( 'feed_ingestion.flows.amazon_digital_services.tasks.garcon_feed_status') def test_fetch_from_drop_location_required_not_found( feed_status_mock, ftp_mock, ftp_conf_mock, mock_config_ads): """Should stop if a file for required country was not found.""" activity_mock = MagicMock() tasks.fetch_from_drop_location( activity_mock, '2018-02-20', 'amazon_digital_services', 's3://somepath/tmp/', 'theorchard', use_s3=None) def copy_side_effect(s3_file_name, **kwargs): return { 'file': s3_file_name, 'status': False } copy_file_from_ftp_to_s3_mock = MagicMock() copy_file_from_ftp_to_s3_mock.side_effect = copy_side_effect ftp_mock.copy_file_from_ftp_to_s3 = copy_file_from_ftp_to_s3_mock activity_mock = MagicMock() result = tasks.fetch_from_drop_location( activity_mock, '2018-02-20', 'amazon_digital_services', 's3://somepath/tmp/', 'theorchard', use_s3=None) assert result == { 'message': 'Not all the required files are present yet', 'missing_files': 's3://somepath/tmp/ZQOR_20180220_Summary_Statement_AT.csv.zip,' 's3://somepath/tmp/ZQOR_20180220_Summary_Statement_US.csv.zip', 'stop': True} @patch('feed_ingestion.flows.amazon_digital_services.tasks.s3') @patch('feed_ingestion.flows.amazon_digital_services.tasks.gzip') @patch('os.remove') @patch('builtins.open') def test_prepare_files_for_ingestion( open_mock, os_remove_mock, gzip_mock, s3_mock, monkeypatch): """Test prepare_files_for_ingestion task.""" activity_mock = MagicMock() source_files_dict = {'files': [ { 'file_name': 'awesome_file.zip', 'found': True, 'file_size': 1000 }, { 'file_name': 'bad_file.zip', 'found': False, 'file_size': -1 } ]} local_temp_path = os.path.realpath(os.getcwd()) + '/' unzip_file_mock = MagicMock(return_value='f/full_path.csv') monkeypatch.setattr(tasks, '_unzip_file', unzip_file_mock) remove_footer_mock = MagicMock() monkeypatch.setattr(tasks, '_remove_footer', remove_footer_mock) tasks.prepare_files_for_ingestion( activity_mock, 'dev-cucumbers', '2021-01-01', source_files_dict, 's3://tmp/', 's3://archive/', 'theorchard') s3_mock.helpers.download.assert_called_with( 's3://tmp/awesome_file.zip', local_temp_path + 'awesome_file.zip', config.expected_bucket_owner ) unzip_file_mock.assert_called_with(local_temp_path + 'awesome_file.zip') remove_footer_mock.assert_called_with('f/full_path.csv', 1) open_mock.assert_called_with('f/full_path.csv', 'rb') gzip_mock.open.assert_called_with( local_temp_path + 'awesome_file.gz', 'wb') s3_mock.helpers.upload_raw_file_to_s3.assert_called_with( local_temp_path + 'awesome_file.gz', 's3://archive/awesome_file.gz', '437795906767' ) os_remove_mock.assert_has_calls([ call('f/full_path.csv'), call(local_temp_path + 'awesome_file.zip'), call(local_temp_path + 'awesome_file.gz')]) def test_notify_missing_files(mock_boto3): """Test notify_missing_files activity.""" source_files_dict = { 'files': [ {'file_name': 'test1.zip', 'required_not_found': 'test1.zip'}, {'file_name': 'test2.zip', 'required_not_found': 'test2.zip'}, {'file_name': 'test3.zip', 'required_not_found': False}]} sub = 'Alert: Amazon Digital Services files are missing' body = ( 'Amazon Digital Services missing files:\n' 'test1.zip\n' 'test2.zip' ) tasks.notify_missing_files(MagicMock(), source_files_dict) assert mock_boto3.client.return_value.send_email.call_args_list == [ mock.call(Destination={ 'ToAddresses': config.EMAIL_LIST}, Message={ 'Subject': {'Data': sub, 'Charset': 'UTF-8'}, 'Body': { 'Text': {'Data': body, 'Charset': 'UTF-8'}}}, Source=config.EMAIL_SOURCE) ]