"""Tests for theatrical ETL SWF tasks.""" from unittest import mock from unittest.mock import Mock from unittest.mock import patch import pytest from flows.theatrical import config from flows.theatrical import status from flows.theatrical import tasks def _find_files_mock(bucket_name, archive_prefix, start_date, end_date): if bucket_name == 'archive': return [ {'filename': 'file2'}, {'filename': 'file3'}] else: return [ {'filename': 'file1'}, {'filename': 'file2'}] @pytest.fixture() def bootstrap_params(): """Set bootstrap params.""" return { 'activity': Mock(), 'correlation_id': 'cid-1234.1', 'date_start': '2016-01-01', 'date_end': '2016-01-01', 'unload_bucket': 'unload', 'unload_prefix': '/', 'archive_bucket': 'archive', 'archive_prefix': '/', 'with_archive': True} @patch('flows.theatrical.tasks.log') @patch('flows.theatrical.tasks.s3') def test_bootstrap_task(s3, mock_log, bootstrap_params): """Test for bootstrap task.""" s3.find_files = Mock(side_effect=_find_files_mock) correlation_id = bootstrap_params['correlation_id'] result = tasks.bootstrap(**bootstrap_params) assert result['source_files'] == [ {'filename': 'file1', 'new': True}, {'filename': 'file2', 'new': True}, {'filename': 'file3'}] mock_log.update_status.assert_called_once_with( correlation_id, status.BOOTSTRAPPED) @patch('flows.theatrical.tasks.log') @patch('flows.theatrical.tasks.s3') def test_bootstrap_task_no_files(s3, mock_log, bootstrap_params): """Test bootstrap task in case source files not found.""" s3.find_files.return_value = [] correlation_id = bootstrap_params['correlation_id'] result = tasks.bootstrap(**bootstrap_params) mock_log.update_status.assert_called_once_with( correlation_id, status.TERMINATED_NO_FILES) assert result['stop'] is True @pytest.fixture() def source_files(): """Set source files fixture.""" return [{'s3_path': '', 'filename': '', 'bucket_name': '', 'date_start': '2016-11-14', 'date_end': '2016-11-17', 'batch_date': '2016-11-19', 'new': True}, {'s3_path': '', 'filename': '', 'bucket_name': '', 'date_start': '2016-11-18', 'date_end': '2016-11-20', 'batch_date': '2016-11-23'}] @pytest.fixture(params=[ ['123456789012', '123456789013'], [] ]) def upcs(request): """Set upcs fixture.""" return request.param @patch('flows.theatrical.tasks.utils') @patch('flows.theatrical.tasks.queries') @patch('flows.theatrical.tasks.datastore') @patch('flows.theatrical.tasks.dtd_utils') @patch('flows.theatrical.tasks.log') @patch('flows.theatrical.tasks.s3') def test_download_to_db_task( s3_mock, log_mock, dtd_utils_mock, datastore_mock, queries_mock, utils_mock, source_files, database_context): """Test of download_to_db_task.""" s3_mock.download_csv.return_value = ([], []) dtd_utils_mock.filter_and_transform.return_value = ([], [], set()) datastore_mock.context = database_context correlation_id = 'cid-1234.1' params = { 'activity': Mock(), 'correlation_id': correlation_id, 'source_files': source_files} tasks.download_to_db(**params) database_context._cursor.execute.assert_any_call('START TRANSACTION;') assert utils_mock.deserialize_date.called assert s3_mock.download_csv.called assert dtd_utils_mock.filter_and_transform.called assert queries_mock.clear_existing_data.called assert queries_mock.insert_raw_data.called assert database_context._cursor.executemany.called log_mock.update_status.assert_called_once_with( correlation_id, status.UNLOADED) @patch('flows.theatrical.tasks.datastore') @patch('flows.theatrical.tasks.dtd_utils') @patch('flows.theatrical.tasks.log') @patch('flows.theatrical.tasks.s3') def test_download_to_db_task_failure( s3_mock, log_mock, dtd_utils_mock, datastore_mock, source_files): """Test of download_to_db_task.""" s3_mock.download_csv.return_value = ([], []) dtd_utils_mock.filter_and_transform.return_value = ([], [], set()) cursor = Mock() cursor.executemany.side_effect = Exception() db_context = Mock() db_context.__enter__ = Mock(return_value=(cursor, Mock())) db_context.__exit__ = Mock() callable_object = Mock() callable_object.return_value = db_context datastore_mock.context = callable_object correlation_id = 'cid-1234.1' params = { 'activity': Mock(), 'correlation_id': correlation_id, 'source_files': source_files} result = tasks.download_to_db(**params) log_mock.update_status.assert_called_once_with( correlation_id, status.TERMINATED_BAD_SOURCE_FILES) assert result['stop'] is True @pytest.fixture def correlation_id(): """Set correlation_id fixture.""" return 'test_correlation_id' @pytest.fixture() def dynamo_source_files(): """Set source files fixture for testing DynamoDB statuses.""" return [{'batch_date': 'test_date'}] @patch('flows.theatrical.tasks.log') @patch('flows.theatrical.tasks.garcon_feed_status') def test_clean_dynamo_status( garcon_feed_status_mock, log_mock, correlation_id, dynamo_source_files): """Test clean_dynamo_status task.""" tasks.clean_dynamo_status(Mock(), correlation_id, dynamo_source_files) assert garcon_feed_status_mock.delete_status.called assert garcon_feed_status_mock.delete_status.call_count == len( dynamo_source_files) log_mock.update_status.assert_called_with( correlation_id, status.DYNAMO_STATUS_CLEANED) @patch('flows.theatrical.tasks.log') @patch('flows.theatrical.tasks.s3') def test_move_source_files_to_archive( s3_mock, log_mock, correlation_id, source_files): """Test move_source_files_to_archive task.""" # test data expected_files_for_move = [] for file in source_files: if file.get('new', False): expected_files_for_move.append( {'old_bucket': file['bucket_name'], 'old_key_path': file['s3_path'], 'new_bucket': config.ARCHIVE_BUCKET, 'new_dir': config.ARCHIVE_PREFIX, 'file_name': file['filename']}) # test function call tasks.move_source_files_to_archive( Mock(), correlation_id, source_files) # assertions s3_mock.move_files.assert_called_once_with(expected_files_for_move) log_mock.update_status.assert_called_with( correlation_id, status.SOURCE_FILES_ARCHIVED) @patch('flows.theatrical.tasks.log') @patch('flows.theatrical.tasks.garcon_feed_status') def test_set_dynamo_status( garcon_feed_status_mock, log_mock, correlation_id, dynamo_source_files): """Test clean_dynamo_status task.""" tasks.set_dynamo_status(Mock(), correlation_id, dynamo_source_files) assert garcon_feed_status_mock.set_overall_status.called assert garcon_feed_status_mock.set_overall_status.call_count == len( dynamo_source_files) log_mock.update_status.assert_called_with( correlation_id, status.DYNAMO_STATUS_UPDATED) @patch('flows.theatrical.tasks.log') @patch('flows.theatrical.tasks.util') @patch('flows.theatrical.tasks.queries') @patch('flows.theatrical.tasks.datastore') def test_create_theatrical_revenue_temp_table( datastore_mock, queries_mock, util_mock, log_mock, correlation_id): """Test create_theatrical_revenue_temp_table task.""" util_mock.correlation_id_hex.return_value = correlation_id table_name = '{table_name}_{correlation_id}'.format( table_name=config.THEATRICAL_REVENUE, correlation_id=correlation_id) create_sql = 'sql text' queries_mock.CREATE_TEMP_TABLE.format.return_value = create_sql response = tasks.create_theatrical_revenue_temp_table( Mock(), correlation_id) util_mock.correlation_id_hex.assert_called_once_with(correlation_id) datastore_mock.execute.assert_called_once_with(create_sql) log_mock.update_status.assert_called_once_with( correlation_id, status.TEMP_TABLE_CREATED) assert response == {'temp_table_name': table_name} @patch('flows.theatrical.tasks.log') @patch('flows.theatrical.tasks.utils') @patch('flows.theatrical.tasks.datastore') @patch('flows.theatrical.tasks.queries') def test_insert_to_theatrical_revenue_temp_table( queries_mock, datastore_mock, utils_mock, log_mock, correlation_id, source_files, upcs, database_context): """Test insert_to_theatrical_revenue_temp_table function.""" datastore_mock.context = database_context test_table_name = 'test_table' insert_sql = ['insert sql'] queries_mock.insert_to_temp_from_raw.return_value = insert_sql tasks.insert_to_theatrical_revenue_temp_table( Mock(), correlation_id, test_table_name, source_files, upcs) queries_mock.insert_to_temp_from_raw.assert_called_once_with( test_table_name, upcs) database_context._cursor.execute.assert_any_call('START TRANSACTION;') database_context._cursor.execute.assert_any_call(insert_sql[0], mock.ANY) log_mock.update_status.assert_called_once_with( correlation_id, status.TEMP_TABLE_INSERTED) @patch('flows.theatrical.tasks.log') @patch('flows.theatrical.tasks.datastore') @patch('flows.theatrical.tasks.queries') def test_insert_new_data( queries_mock, datastore_mock, log_mock, database_context, source_files): """Test to insert_new_data task.""" datastore_mock.context = database_context test_table_name = 'test_table' queries_mock.COPY_TO_PROD.format.return_value = 'copy_query' queries_mock.DROP_TABLE.format.return_value = 'drop_query' tasks.insert_new_data( Mock(), correlation_id, test_table_name, source_files, []) database_context._cursor.execute.assert_any_call('START TRANSACTION;') database_context._cursor.execute.assert_any_call('copy_query') database_context._cursor.execute.assert_any_call('drop_query') assert queries_mock.clear_existing_data.called log_mock.update_status.assert_called_once_with( correlation_id, status.TEMP_TABLE_MADE_LIVE) @patch('flows.theatrical.tasks.log') def test_set_final_status(log_mock): """Test to set_final_status task.""" tasks.set_final_status(Mock(), correlation_id) log_mock.update_status.assert_called_once_with( correlation_id, status.INGESTED) @patch('flows.theatrical.tasks.utils') def test_send_notification(utils_mock, upcs): """Test to set_final_status task with db param UPCs.""" correlation_id = '0123-4567-8910-1112' upcs = 'upcs' tasks.send_notification(Mock(), correlation_id, upcs) utils_mock.send_success_notification.assert_called_once_with( correlation_id, upcs)