"""Unit tests for Line tasks workflow.""" from unittest.mock import MagicMock from unittest.mock import patch from boto3.exceptions import S3UploadFailedError from botocore.exceptions import ClientError from garcon_contrib.dynamo_feed_status import garcon_feed_status import pytest from requests import ConnectionError # noqa:A004 from requests import HTTPError from feed_ingestion.flows.line import config from feed_ingestion.flows.line import tasks _date = '2021-01-28' @pytest.fixture def feed_name(): """Feed name value.""" return 'line' @pytest.fixture def source_files_dict(): """Source files value.""" return { 'place': 'LMA_bl00000000029c7cd8_20210128_place.csv', 'pref': 'LMA_bl00000000029c7cd8_20210128_pref.csv', 'sex_age': 'LMA_bl00000000029c7cd8_20210128_sex_age.csv', 'sex_age_hour': 'LMA_bl00000000029c7cd8_20210128_sex_age_hour.csv', 'sex_age_region': 'LMA_bl00000000029c7cd8_20210128_sex_age_region.csv', 'sex_age_type': 'LMA_bl00000000029c7cd8_20210128_sex_age_type.csv', 'type': 'LMA_bl00000000029c7cd8_20210128_type.csv' } @pytest.fixture def expected_bootstrap_response(): """Response for bootstrap task.""" return { 'feed_name': 'line_theorchard', 'reports': [ 'place', 'type', 'pref', 'sex_age_type', 'sex_age_region', 'sex_age', 'sex_age_hour'], 'date': _date, 'licensor': 'theorchard', 'zipped_key_name': 'LINE/zipped_archives/2021-01-28/LMA_bl00000000029c7cd8_20210128' '.zip', 'zipped_filename': 'LMA_bl00000000029c7cd8_20210128.zip', 'zipped_path': 'LINE/zipped_archives/2021-01-28/', 'archive_path': 'LINE/archives/2021-01-28/', 's3_archive_path': 's3://dev-cucumbers/LINE/archives/2021-01-28/', 'source_files_dict': { 'place': 'LMA_bl00000000029c7cd8_20210128_place.csv', 'pref': 'LMA_bl00000000029c7cd8_20210128_pref.csv', 'sex_age': 'LMA_bl00000000029c7cd8_20210128_sex_age.csv', 'sex_age_hour': 'LMA_bl00000000029c7cd8_20210128_sex_age_hour.csv', 'sex_age_region': 'LMA_bl00000000029c7cd8_20210128_sex_age_region.csv', 'sex_age_type': 'LMA_bl00000000029c7cd8_20210128_sex_age_type.csv', 'type': 'LMA_bl00000000029c7cd8_20210128_type.csv' }, 'kwargs': {'licensor': 'theorchard'}, } @patch.object(tasks, 'garcon_feed_status') def test_bootstrap(garcon_feed_status_mock, expected_bootstrap_response): """Test bootstrap.""" result = tasks.bootstrap( MagicMock(), _date, licensor='theorchard', reload=False) assert result == expected_bootstrap_response @pytest.fixture def mock_task_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 mock_set_overall_status(): """Yield overall status.""" overall_status_path = ( 'feed_ingestion.flows.line.tasks.garcon_feed_status.' 'set_overall_status') with patch(overall_status_path) as overall_status: yield overall_status @pytest.fixture def mock_delete_status(): """Yield delete overall status.""" path = 'feed_ingestion.flows.line.tasks.garcon_feed_status.delete_status' with patch(path) as delete_status: yield delete_status @pytest.fixture def mock_set_missing_files(): """Yield set_missing_files.""" path = ( 'feed_ingestion.flows.line.tasks.garcon_feed_status.' 'set_missing_files') with patch(path) as mock_set_missing: yield mock_set_missing @pytest.fixture def mock_s3_tasks(): """Mock s3_tasks.""" path = 'feed_ingestion.flows.line.tasks.s3_tasks' with patch(path) as mock_s3: yield mock_s3 class TestGrabDropFiles(object): """Test grab_drop_files.""" @pytest.fixture def context_grab_drop_files(self, monkeypatch): """Return context for grab_drop_files.""" context = dict( activity=MagicMock(), feed_name=config.feed_name, date=_date, licensor='theorchard', zipped_filename='zipped_filename', zipped_path='zipped_path') return context @pytest.fixture def mock_send_error_or_warning(self): """Mock function send_error_or_warning.""" path = 'feed_ingestion.flows.spotify.tasks.send_error_or_warning' with patch(path) as send_error_or_warning: yield send_error_or_warning @pytest.fixture def mock_line_api(self): """Line API Wrapper fixture.""" la_path = ( 'feed_ingestion.flows.line.' 'tasks.LineAPI') with patch(la_path) as line_api: api_instance = line_api.return_value api_instance.request_to_file = MagicMock() yield api_instance @pytest.fixture def mock_line_api_fail(self): """Line API Wrapper failing fixture.""" la_path = ( 'feed_ingestion.flows.line.' 'tasks.LineAPI') mock_response = MagicMock() mock_response.status_code = 500 err = HTTPError(response=mock_response) with patch(la_path) as line_api: api_instance = line_api.return_value api_instance.request_to_file = ( MagicMock(side_effect=err)) yield api_instance @pytest.fixture def mock_line_api_not_available(self): """Line API Wrapper resource not available fixture.""" la_path = ( 'feed_ingestion.flows.line.' 'tasks.LineAPI') mock_response = MagicMock() mock_response.status_code = 404 err = HTTPError(response=mock_response) with patch(la_path) as line_api: api_instance = line_api.return_value api_instance.request_to_file = ( MagicMock(side_effect=err)) yield api_instance @pytest.fixture def mock_line_api_connection_error(self): """Line API Wrapper resource connection error.""" la_path = 'feed_ingestion.flows.line.tasks.LineAPI' mock_response = MagicMock() mock_response.status_code = 404 err = ConnectionError(response=mock_response) with patch(la_path) as line_api: api_instance = line_api.return_value api_instance.request_to_file = ( MagicMock(side_effect=err)) yield api_instance @pytest.fixture def mock_boto3_fail(self): """Mock boto3.""" boto3_path = 'feed_ingestion.flows.line.tasks.boto3' with patch(boto3_path) as boto3: boto3.client.side_effect = S3UploadFailedError('Failed to upload') yield boto3 @pytest.fixture def mock_remove_files_from_path(self): """Return mock remove_files_from_path.""" remove_files_path = ( 'feed_ingestion.flows.line.tasks.remove_files_from_path') with patch(remove_files_path) as remove_files: yield remove_files @pytest.fixture def grab_drop_files( self, context_grab_drop_files, mock_boto3, mock_task_status, mock_set_overall_status): """Run grab_drop_files.""" tasks.grab_drop_files(**context_grab_drop_files) def test_clear_s3( self, mock_line_api, mock_boto3, mock_remove_files_from_path, grab_drop_files): """Should clear archive destination folder before uploading to S3.""" assert mock_remove_files_from_path.called def test_request_to_file_api( self, mock_line_api, mock_remove_files_from_path, grab_drop_files): """Should request tracks file from LineAPI.""" assert mock_line_api.request_to_file.called def test_upload_on_s3( self, mock_line_api, mock_boto3, mock_remove_files_from_path, grab_drop_files): """Should upload file to S3.""" assert mock_boto3.upload_file.called def test_set_overall_status( self, mock_line_api, mock_task_status, mock_boto3, mock_remove_files_from_path, grab_drop_files): """Should set status after successful upload.""" assert mock_task_status.mark_completed_task.called def test_api_failure( self, context_grab_drop_files, mock_line_api_fail, mock_task_status, mock_set_overall_status, mock_boto3, mock_remove_files_from_path, mock_send_error_or_warning): """Test SpotifyAPI failure.""" tasks.grab_drop_files(**context_grab_drop_files) mock_set_overall_status.assert_called_with( config.feed_name, _date, garcon_feed_status.STATUS_NOT_AVAILABLE) mock_task_status.mark_completed_task.assert_not_called() def test_api_resource_not_available( self, context_grab_drop_files, mock_line_api_not_available, mock_task_status, mock_set_overall_status, mock_boto3, mock_remove_files_from_path, mock_send_error_or_warning): """Test LineAPI resource is not available.""" res = tasks.grab_drop_files(**context_grab_drop_files) assert res == {'stop': True} mock_set_overall_status.assert_called_with( config.feed_name, _date, garcon_feed_status.STATUS_NOT_AVAILABLE) mock_task_status.mark_completed_task.assert_not_called() mock_send_error_or_warning.assert_not_called() def test_api_resource_connection_error( self, context_grab_drop_files, mock_line_api_connection_error, mock_task_status, mock_set_overall_status, mock_boto3, mock_remove_files_from_path, mock_send_error_or_warning): """Test LineAPI resource returns connection error.""" res = tasks.grab_drop_files(**context_grab_drop_files) assert res == {'stop': True} mock_set_overall_status.assert_called_with( config.feed_name, _date, garcon_feed_status.STATUS_NOT_AVAILABLE) mock_task_status.mark_completed_task.assert_not_called() def test_s3_failure( self, context_grab_drop_files, mock_line_api, mock_task_status, mock_set_overall_status, mock_boto3_fail, mock_remove_files_from_path): """Test S3 failure.""" with pytest.raises(S3UploadFailedError): tasks.grab_drop_files(**context_grab_drop_files) mock_set_overall_status.assert_called_with( config.feed_name, _date, garcon_feed_status.STATUS_NOT_AVAILABLE) mock_task_status.mark_completed_task.assert_not_called() @pytest.fixture def mock_boto3(): """Mock boto3.""" boto3_path = 'feed_ingestion.flows.line.tasks.boto3' with patch(boto3_path) as boto3: mock_client = MagicMock() boto3.client.return_value = mock_client yield mock_client @pytest.fixture def mock_shutil(): """Mock shutil.""" path = 'feed_ingestion.flows.line.tasks.shutil' with patch(path) as shutil: yield shutil @pytest.fixture def mock_zipfile(): """Mock zipfile module.""" path = ( 'feed_ingestion.flows.line.tasks.zipfile.ZipFile') with patch(path) as mock_zipfile: mock_zipfile_obj = MagicMock() mock_zipfile.return_value.__enter__.return_value = ( mock_zipfile_obj) yield mock_zipfile_obj @pytest.fixture def mock_os(): """Mock os.""" path = 'feed_ingestion.flows.line.tasks.os.makedirs' with patch(path) as mock_makedirs: yield mock_makedirs def test_unzip_files_no_files_in_archive( mock_task_status, mock_zipfile, mock_boto3, mock_set_overall_status, mock_delete_status, mock_shutil, mock_os): """Test unzip_files if there is no files in an archive.""" mock_zipfile.namelist.return_value = [''] result = tasks.unzip_files( activity=MagicMock(), feed_name='line', date=_date, archive_path='archive_path', zipped_key_name='zipped_key_name', zipped_filename='zipped_filename.zip' ) assert result == {'stop': True} mock_delete_status.assert_called_once_with('line', _date) mock_set_overall_status.assert_called_once_with( 'line', _date, garcon_feed_status.STATUS_NOT_INGESTED) mock_boto3.download_file.assert_called_once_with( 'dev-cucumbers', 'zipped_key_name', './file_stage_2021-01-28/zipped_filename.zip') mock_shutil.rmtree.assert_called_once_with('./file_stage_2021-01-28') def test_unzip_files( mock_task_status, mock_zipfile, mock_boto3, mock_set_overall_status, mock_delete_status, mock_shutil, mock_os): """Test unzip_files.""" mock_zipfile.namelist.return_value = [ 'LMA_bl00000000029c7cd8_20210128_place.csv', 'LMA_bl00000000029c7cd8_20210128_pref.csv', 'LMA_bl00000000029c7cd8_20210128_sex_age_hour.csv', 'LMA_bl00000000029c7cd8_20210128_sex_age_region.csv', 'LMA_bl00000000029c7cd8_20210128_sex_age_type.csv', 'LMA_bl00000000029c7cd8_20210128_sex_age.csv', 'LMA_bl00000000029c7cd8_20210128_type.csv', ] tasks.unzip_files( activity=MagicMock(), feed_name='line', date=_date, archive_path='archive_path', zipped_key_name='zipped_key_name', zipped_filename='zipped_filename.zip' ) mock_boto3.download_file.assert_called_once_with( 'dev-cucumbers', 'zipped_key_name', './file_stage_2021-01-28/zipped_filename.zip') assert mock_boto3.upload_file.call_count == 7 mock_delete_status.assert_called_with('line', '2021-01-28') mock_set_overall_status.assert_called_with( 'line', '2021-01-28', 'NOT_INGESTED') mock_shutil.rmtree.assert_called_once_with('./file_stage_2021-01-28') def test_unzip_files_failed_while_uploading_files( mock_task_status, mock_zipfile, mock_boto3, mock_set_overall_status, mock_delete_status, mock_shutil, mock_os): """Test unzip_files if copy files on s3 failed.""" mock_zipfile.namelist.return_value = [ 'LMA_bl00000000029c7cd8_20210128_place.csv', 'LMA_bl00000000029c7cd8_20210128_pref.csv', 'LMA_bl00000000029c7cd8_20210128_sex_age_hour.csv', 'LMA_bl00000000029c7cd8_20210128_sex_age_region.csv', 'LMA_bl00000000029c7cd8_20210128_sex_age_type.csv', 'LMA_bl00000000029c7cd8_20210128_sex_age.csv', 'LMA_bl00000000029c7cd8_20210128_type.csv', ] mock_boto3.upload_file.side_effect = ClientError( {'Error': {'Code': '404', 'Message': 'Not found'}}, 'Not found') with pytest.raises(ClientError): tasks.unzip_files( activity=MagicMock(), feed_name='line', date=_date, archive_path='archive_path', zipped_key_name='zipped_key_name', zipped_filename='zipped_filename.zip' ) mock_boto3.download_file.assert_called_once_with( 'dev-cucumbers', 'zipped_key_name', './file_stage_2021-01-28/zipped_filename.zip') assert mock_boto3.upload_file.call_count == 1 mock_boto3.download_file.assert_called_once_with( 'dev-cucumbers', 'zipped_key_name', './file_stage_2021-01-28/zipped_filename.zip') mock_delete_status.assert_called_once_with('line', _date) mock_set_overall_status.assert_called_once_with( 'line', _date, garcon_feed_status.STATUS_NOT_INGESTED) mock_shutil.rmtree.assert_called_once_with('./file_stage_2021-01-28') @patch('feed_ingestion.flows.line.tasks.task_status') @patch('feed_ingestion.flows.line.tasks.LINE') @patch('feed_ingestion.flows.line.tasks.get_sf_config') def test_create_temp_staging_raw_table( mock_get_sf_config, mock_sf_executor_class, mock_task_status, feed_name): """Test create temp staging raw table.""" mock_sf_config = {'db': 'DB', 'schema': 'SCHEMA'} mock_get_sf_config.return_value = mock_sf_config mock_sf_executor = MagicMock() mock_sf_executor_class.return_value.__enter__.return_value = \ mock_sf_executor mock_sf_executor.create_temp_staging_raw_table.return_value = None mock_task_status.is_completed_task.return_value = False report = 'place' date = '2021-01-28' temp_table = 'temp_staging_raw_line_place_2021_01_28' tasks.create_temp_staging_raw_table(MagicMock(), feed_name, date, report, temp_table) mock_get_sf_config.assert_called_once_with('swf-line') mock_sf_executor.create_temp_staging_raw_table.assert_called_once_with( temp_table, report_type='place') @patch('feed_ingestion.flows.line.tasks.task_status') @patch('feed_ingestion.flows.line.tasks.LINE') @patch('feed_ingestion.flows.line.tasks.get_sf_config') def test_load_temp_staging_raw_table( mock_get_sf_config, mock_sf_executor_class, mock_task_status, feed_name, source_files_dict): """Test load temp staging raw table.""" mock_sf_config = {'db': 'DB', 'schema': 'SCHEMA'} mock_get_sf_config.return_value = mock_sf_config mock_sf_executor = MagicMock() mock_sf_executor_class.return_value.__enter__.return_value = \ mock_sf_executor mock_sf_executor.load_temp_staging_raw_table.return_value = None mock_task_status.is_completed_task.return_value = False report = 'place' s3_dir_path = 'http://some_path/' file_path = '{}{}'.format(s3_dir_path, source_files_dict[report]) date = '2021-01-28' temp_table = 'temp_staging_raw_line_place_2021_01_28' tasks.load_temp_staging_raw_table( MagicMock(), feed_name, date, s3_dir_path, source_files_dict, report, temp_table) mock_get_sf_config.assert_called_once_with('swf-line') mock_sf_executor.load_temp_staging_raw_table.assert_called_once_with( temp_table, None, file_path, error_limit=1) @patch('feed_ingestion.flows.line.tasks.task_status') @patch('feed_ingestion.flows.line.tasks.LINE') @patch('feed_ingestion.flows.line.tasks.get_sf_config') def test_load_staging_raw_table( mock_get_sf_config, mock_sf_executor_class, mock_task_status, feed_name): """Test load staging raw table.""" mock_sf_config = {'db': 'DB', 'schema': 'SCHEMA'} mock_get_sf_config.return_value = mock_sf_config mock_sf_executor = MagicMock() mock_sf_executor_class.return_value.__enter__.return_value = \ mock_sf_executor mock_sf_executor.load_staging_raw_table.return_value = None mock_sf_executor.clean_staging_raw_table.return_value = None mock_task_status.is_completed_task.return_value = False report = 'place' date = '2021-01-28' staging_raw_table = 'staging_raw_line_place' temp_table_name = 'temp_staging_raw_line_{}_2021_01_28'.format(report) kwargs = dict(licensor='somelicensor') tasks.load_staging_raw_table( MagicMock(), feed_name, date, temp_table_name, report, staging_raw_table, kwargs) mock_get_sf_config.assert_called_once_with('swf-line') mock_sf_executor.clean_staging_raw_table.assert_called_once_with( '2021-01-28', 'staging_raw_line_place', licensor='somelicensor') mock_sf_executor.load_staging_raw_table.assert_called_once_with( temp_table_name, 'staging_raw_line_place', 'place', '2021-01-28', licensor='somelicensor') @patch('feed_ingestion.flows.line.tasks.LINE') @patch('feed_ingestion.flows.line.tasks.get_sf_config') def test_delete_tmp_table( mock_get_sf_config, mock_sf_executor_class, feed_name): """Test delete staging fact and temp staging raw tables.""" mock_sf_config = {'db': 'DB', 'schema': 'SCHEMA'} mock_get_sf_config.return_value = mock_sf_config mock_sf_executor = MagicMock() mock_sf_executor_class.return_value.__enter__.return_value = \ mock_sf_executor mock_sf_executor.drop_temp_staging_raw_table.return_value = None tasks.delete_temp_staging_raw_table(MagicMock(), 'temp_table_name') mock_get_sf_config.assert_called_once_with('swf-line') mock_sf_executor.drop_temp_staging_raw_table.assert_called_once_with( 'temp_table_name')