"""Unit tests for tasks for Pandora Analytics Ingestion Workflow.""" from importlib import reload from unittest.mock import MagicMock from unittest.mock import patch from freezegun import freeze_time from moto import mock_aws import pytest from feed_ingestion import tasks as common_tasks from feed_ingestion.flows.pandora import config from feed_ingestion.flows.pandora import tasks _mock_config = { 'feed_name': 'pandora', 'secrets_path': config.secrets_path, 'pandora_streams_file': config.pandora_streams_file, 'pandora_metadata_file': config.pandora_metadata_file, 'countries': config.countries, 's3': { 'drop': { 'theorchard': 's3://dev-feed-drop/data/{YYYY}/{MM}/{DD}/', 'sme': ( 's3://sme-ca-prod-partners/pandora/in/pandora-streams/' '{YYYY}/{MM}/{DD}/')}, 'archive': { 'theorchard': 's3://dev-cucumbers/Pandora/archives/{date}/', 'sme': 's3://dev-cucumbers/Pandora/archives_sme/{date}/'}}, 'licensors': { 'theorchard': 'orchard', 'sme': 'sme'}, 'snowflake_table_names': config.snowflake_table_names, 'default_snowflake_error_limit': 5, 'jenkins_config': config.jenkins_config, } @pytest.fixture(autouse=True) def mock_check_status(request, monkeypatch): """Mock check_status decorator. Mock so that the decorated function is kept intact and always get called. """ mock = MagicMock() # just return the function without any modifications mock.return_value = lambda f: f monkeypatch.setattr(common_tasks, 'check_status', value=mock) reload(tasks) # redecorate tasks yield monkeypatch.undo() reload(tasks) @pytest.fixture def mock_executor_context(): """Yield executor context.""" sf_executor_class_path = 'feed_ingestion.flows.pandora.tasks.Pandora' 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_registered_executors(): """Yield executor context.""" sf_executor_class_path = ( 'feed_ingestion.flows.pandora.tasks.registered_executors') with patch(sf_executor_class_path) as sf_executor: mock_executor_context = MagicMock() sf_executor.get.return_value = mock_executor_context yield mock_executor_context.return_value.__enter__.return_value @patch('feed_ingestion.flows.pandora.tasks.task_status') def _run_completed_task(task, context, mock_task_status): """Check that task runs correctly when status is marked as completed. Args: task (callable): Garcon task to call. context (dict): Context to pass to the task. mock_task_status (MagicMock): Mock object for task status. Returns: return_value (dict or None): Return value of task. """ mock_task_status.is_completed_task.return_value = True return_value = task(MagicMock(), **context) if context.get('date'): mock_task_status.is_completed_task.assert_called_with( context['feed_name'], context['date'], 'staging_raw_table_tasks') assert not mock_task_status.mark_completed_task.called return return_value @patch('feed_ingestion.flows.pandora.tasks.task_status') def _run_task(task, context, mock_task_status): """Check that task runs correctly when status is marked as not completed. Args: task (callable): Garcon task to call. context (dict): Context to pass to the task. mock_task_status (MagicMock): Mock object for task status. Returns: return_value (dict or None): Return value of task. """ mock_task_status.is_completed_task.return_value = False return_value = task(MagicMock(), **context) if context.get('date'): mock_task_status.is_completed_task.assert_called_with( context['feed_name'], context['date'], 'staging_raw_table_tasks') return return_value @pytest.fixture def mock_delete_status(): """Yield delete status.""" delete_status_path = ( 'feed_ingestion.flows.pandora.tasks.garcon_feed_status.' 'delete_status') with patch(delete_status_path) as delete_status: yield delete_status @pytest.fixture def mock_get_overall_status(): """Yield overall status.""" overall_status_path = ( 'feed_ingestion.flows.pandora.tasks.garcon_feed_status.' 'get_overall_status') with patch(overall_status_path) as overall_status: yield overall_status @pytest.mark.parametrize('context, expected', [ ( # default date { 'date': None, 'reload': False, 'licensor': 'theorchard', 'snowflake_error_on_column_count_mismatch': 'False', }, { 'date': '2017-10-01', 'snowflake_error_limit': 5, 'snowflake_error_on_column_count_mismatch': 'False', 'delete_status_params': None, } ), ( # specify explicit date { 'date': '2020-12-31', 'reload': False, 'licensor': 'theorchard', 'snowflake_error_on_column_count_mismatch': 'falSe', }, { 'date': '2020-12-31', 'snowflake_error_limit': 5, 'snowflake_error_on_column_count_mismatch': 'falSe', 'delete_status_params': None, } ), ( # set reload { 'date': None, 'reload': 'True', 'licensor': 'theorchard', 'snowflake_error_on_column_count_mismatch': 'FALSE', }, { 'date': '2017-10-01', 'snowflake_error_limit': 5, 'snowflake_error_on_column_count_mismatch': 'FALSE', 'delete_status_params': ['pandora_theorchard', '2017-10-01'], } ), ( # specify snowflake_error_limit { 'date': None, 'reload': False, 'licensor': 'theorchard', 'snowflake_error_on_column_count_mismatch': None, 'snowflake_error_limit': 105, }, { 'date': '2017-10-01', 'snowflake_error_limit': 105, 'snowflake_error_on_column_count_mismatch': 'true', 'delete_status_params': None, } ), ( # specify incorrect snowflake_error_limit { 'date': None, 'reload': False, 'licensor': 'theorchard', 'snowflake_error_on_column_count_mismatch': 'false', 'snowflake_error_limit': 'DROP FROM *', }, { 'date': '2017-10-01', 'snowflake_error_limit': ValueError, 'snowflake_error_on_column_count_mismatch': 'false', 'delete_status_params': None, } ), ]) @freeze_time('2017-10-01') @patch('feed_ingestion.flows.pandora.tasks.config', **_mock_config) @patch('feed_ingestion.flows.pandora.tasks.task_status.is_completed_task') def test_bootstrap_theorchard( is_completed_task, mock_config, mock_get_overall_status, mock_delete_status, context, expected): """Check that bootstrap returns expected results.""" expected_date = expected['date'] # next two lines are a piece of code taken from # pandora/tasks.py to check the expected result expected_snowflake_mismatch_error = \ context['snowflake_error_on_column_count_mismatch'] or 'true' assert expected['snowflake_error_on_column_count_mismatch'] == \ expected_snowflake_mismatch_error expected_date_slashes = expected_date.replace('-', '/') expected_date_slim = expected_date.replace('-', '') expected_response = { 'feed_name': 'pandora_theorchard', 'secrets_path': 'pandora', 'replace_archive_files': False, 'date': expected_date, 'processed_datetime': '2017-10-01T00:00:00', 'date_as_in_uuid': expected_date, 'drop_bucket': f's3://dev-feed-drop/data/{expected_date_slashes}/', 'archive_bucket': f's3://dev-cucumbers/Pandora/archives/{expected_date}/', 'expected_files': [ f'orchard_metadata_{expected_date}.txt.bz2', f'orchard_US_{expected_date}.txt.bz2'], 'snowflake_error_on_column_count_mismatch': expected_snowflake_mismatch_error, 'temp_staging_raw_tables': { 'metadata': { 'temp_table_name': f'pandora_metadata_theorchard_{expected_date_slim}', 'temp_table_s3_full_path': f's3://dev-cucumbers/Pandora/archives/{expected_date}/' f'orchard_metadata_{expected_date}.txt.bz2'}, 'US': { 'temp_table_name': f'pandora_streams_theorchard_US_{expected_date_slim}', 'temp_table_s3_full_path': f's3://dev-cucumbers/Pandora/archives/{expected_date}/' f'orchard_US_{expected_date}.txt.bz2'}}, 'licensor': 'theorchard', 'snowflake_error_limit': expected['snowflake_error_limit'], 'jenkins_config': config.jenkins_config } activity_mock = MagicMock() if expected['snowflake_error_limit'] == ValueError: with pytest.raises(ValueError): tasks.bootstrap(activity_mock, **context) else: reply = tasks.bootstrap(activity_mock, **context) assert reply == expected_response delete_status_params = expected['delete_status_params'] if not delete_status_params: mock_delete_status.assert_not_called() else: mock_delete_status.assert_called_with(*delete_status_params) @pytest.mark.parametrize('context, expected', [ ( # default date { 'date': None, 'reload': False, 'licensor': 'sme', 'snowflake_error_on_column_count_mismatch': 'true' }, { 'date': '2017-10-01', 'snowflake_error_limit': 5, 'delete_status_params': None, 'snowflake_error_on_column_count_mismatch': 'true' } ), ( # specify explicit date { 'date': '2020-12-30', 'reload': False, 'licensor': 'sme', 'snowflake_error_on_column_count_mismatch': 'True' }, { 'date': '2020-12-30', 'snowflake_error_limit': 5, 'delete_status_params': None, 'snowflake_error_on_column_count_mismatch': 'True' } ), ( # set reload { 'date': None, 'reload': 'True', 'licensor': 'sme', 'snowflake_error_on_column_count_mismatch': None }, { 'date': '2017-10-01', 'snowflake_error_limit': 5, 'delete_status_params': ['pandora_sme', '2017-10-01'], 'snowflake_error_on_column_count_mismatch': 'true' } ), ( # specify snowflake_error_limit { 'date': None, 'reload': False, 'licensor': 'sme', 'snowflake_error_limit': 404, 'snowflake_error_on_column_count_mismatch': 'anything' }, { 'date': '2017-10-01', 'snowflake_error_limit': 404, 'delete_status_params': None, 'snowflake_error_on_column_count_mismatch': 'anything' } ), ( # specify incorrect snowflake_error_limit { 'date': None, 'reload': False, 'licensor': 'sme', 'snowflake_error_limit': '20', 'snowflake_error_on_column_count_mismatch': 'true' }, { 'date': '2017-10-01', 'snowflake_error_limit': ValueError, 'delete_status_params': None, 'snowflake_error_on_column_count_mismatch': 'true' } ), ]) @freeze_time('2017-10-01') @patch('feed_ingestion.flows.pandora.tasks.config', **_mock_config) @patch('feed_ingestion.flows.pandora.tasks.task_status.is_completed_task') def test_bootstrap_sme( is_completed_task, mock_config, mock_get_overall_status, mock_delete_status, context, expected): """Check that bootstrap returns expected results.""" expected_date = expected['date'] # next two lines are a piece of code taken from # pandora/tasks.py to check the expected result expected_snowflake_mismatch_error = \ context['snowflake_error_on_column_count_mismatch'] or 'true' assert expected['snowflake_error_on_column_count_mismatch'] == \ expected_snowflake_mismatch_error expected_date_slashes = expected_date.replace('-', '/') expected_date_slim = expected_date.replace('-', '') expected_response = { 'feed_name': 'pandora_sme', 'secrets_path': 'pandora', 'replace_archive_files': False, 'date': expected_date, 'processed_datetime': '2017-10-01T00:00:00', 'date_as_in_uuid': f'{expected_date}', 'drop_bucket': 's3://sme-ca-prod-partners/pandora/in/' f'pandora-streams/{expected_date_slashes}/', 'archive_bucket': f's3://dev-cucumbers/Pandora/archives_sme/{expected_date}/', 'expected_files': [ f'sme_metadata_{expected_date}.txt.bz2', f'sme_US_{expected_date}.txt.bz2'], 'snowflake_error_on_column_count_mismatch': expected_snowflake_mismatch_error, 'temp_staging_raw_tables': { 'metadata': { 'temp_table_name': f'pandora_metadata_sme_{expected_date_slim}', 'temp_table_s3_full_path': f's3://dev-cucumbers/Pandora/archives_sme/{expected_date}/' f'sme_metadata_{expected_date}.txt.bz2'}, 'US': { 'temp_table_name': f'pandora_streams_sme_US_{expected_date_slim}', 'temp_table_s3_full_path': f's3://dev-cucumbers/Pandora/archives_sme/{expected_date}/' f'sme_US_{expected_date}.txt.bz2'}}, 'licensor': 'sme', 'snowflake_error_limit': expected['snowflake_error_limit'], 'jenkins_config': config.jenkins_config } activity_mock = MagicMock() if expected['snowflake_error_limit'] == ValueError: with pytest.raises(ValueError): tasks.bootstrap(activity_mock, **context) else: reply = tasks.bootstrap(activity_mock, **context) assert reply == expected_response delete_status_params = expected['delete_status_params'] if not delete_status_params: mock_delete_status.assert_not_called() else: mock_delete_status.assert_called_with(*delete_status_params) @patch('feed_ingestion.flows.pandora.tasks.task_status.is_completed_task') def test_bootstrap_replace_archive_files( is_completed_task, mock_get_overall_status): """Check that archive files are replaced if status is not DOWNLOADED.""" is_completed_task.return_value = False response = tasks.bootstrap( MagicMock(), date='2016-01-01', reload=False, licensor='theorchard') assert response['replace_archive_files'] is_completed_task.return_value = True response = tasks.bootstrap( MagicMock(), date='2016-01-01', reload=False, licensor='theorchard') assert not response['replace_archive_files'] @patch('feed_ingestion.util.task_status.is_completed_task') @mock_aws def test_clean_staging_raw_table_theorchard( is_completed_task, mock_registered_executors): """Test clean_staging_raw_table task.""" is_completed_task.return_value = False _run_task(tasks.clean_staging_raw_table, { 'date': '2017-01-01', 'feed_name': 'pandora_theorchard'}) mock_registered_executors.clean_staging_raw_table.assert_called_with( config.snowflake_table_names['staging_raw'], '2017-01-01') mock_registered_executors.reset_mock() _run_completed_task(tasks.clean_staging_raw_table, { 'date': '2017-01-01', 'feed_name': 'pandora_theorchard'}) assert mock_registered_executors.call_count == 0 @patch('feed_ingestion.util.task_status.is_completed_task') @mock_aws def test_clean_staging_raw_table_sme( is_completed_task, mock_registered_executors): """Test clean_staging_raw_table task.""" is_completed_task.return_value = False _run_task(tasks.clean_staging_raw_table, { 'date': '2017-01-01', 'feed_name': 'pandora_sme'}) mock_registered_executors.clean_staging_raw_table.assert_called_with( config.snowflake_table_names['staging_raw'], '2017-01-01') mock_registered_executors.reset_mock() _run_completed_task(tasks.clean_staging_raw_table, { 'date': '2017-01-01', 'feed_name': 'pandora_sme'}) assert mock_registered_executors.call_count == 0 def test_create_temp_staging_raw_table_theorchard(mock_registered_executors): """Test create_temp_staging_raw_table task.""" _run_task(tasks.create_temp_staging_raw_table, { 'date': '2017-01-01', 'temp_table_name': 'streams_test_table', 'feed_name': 'pandora_theorchard'}) (mock_registered_executors.create_streams_temp_staging_raw_table. assert_called_with('streams_test_table', date_for_sqlloader='2017-01-01')) mock_registered_executors.reset_mock() _run_task(tasks.create_temp_staging_raw_table, { 'date': '2017-01-01', 'temp_table_name': 'metadata_test_table', 'feed_name': 'pandora_theorchard'}) (mock_registered_executors.create_metadata_temp_staging_raw_table. assert_called_with('metadata_test_table', date_for_sqlloader='2017-01-01')) mock_registered_executors.reset_mock() _run_completed_task(tasks.create_temp_staging_raw_table, { 'date': '2017-01-01', 'temp_table_name': 'metadata_test_table', 'feed_name': 'pandora_theorchard'}) assert mock_registered_executors.call_count == 0 def test_create_temp_staging_raw_table_sme(mock_registered_executors): """Test create_temp_staging_raw_table task.""" _run_task(tasks.create_temp_staging_raw_table, { 'date': '2017-01-01', 'temp_table_name': 'streams_test_table', 'feed_name': 'pandora_sme'}) (mock_registered_executors.create_streams_temp_staging_raw_table. assert_called_with('streams_test_table', date_for_sqlloader='2017-01-01')) mock_registered_executors.reset_mock() _run_task(tasks.create_temp_staging_raw_table, { 'date': '2017-01-01', 'temp_table_name': 'metadata_test_table', 'feed_name': 'pandora_sme'}) (mock_registered_executors.create_metadata_temp_staging_raw_table. assert_called_with('metadata_test_table', date_for_sqlloader='2017-01-01')) mock_registered_executors.reset_mock() _run_completed_task(tasks.create_temp_staging_raw_table, { 'date': '2017-01-01', 'temp_table_name': 'metadata_test_table', 'feed_name': 'pandora_sme'}) assert mock_registered_executors.call_count == 0 def test_load_staging_raw_table_theorchard(mock_registered_executors): """Test load_staging_raw_table task.""" _run_task(tasks.load_staging_raw_table, { 'date': '2017-01-01', 'processed_datetime': '2017-10-01T00:00:00', 'filename': 'testfilename', 'temp_streams_table': 'streams_test_table', 'temp_metadata_table': 'metadata_test_table', 'feed_name': 'pandora_theorchard'}) (mock_registered_executors.load_staging_raw_table.assert_called_with( '2017-01-01', '2017-10-01T00:00:00', 'testfilename', config.snowflake_table_names['staging_raw'], 'streams_test_table', 'metadata_test_table', date_for_sqlloader='2017-01-01')) mock_registered_executors.reset_mock() _run_completed_task(tasks.load_staging_raw_table, { 'date': '2017-01-01', 'processed_datetime': '2017-10-01T00:00:00', 'filename': 'testfilename', 'temp_streams_table': 'streams_test_table', 'temp_metadata_table': 'metadata_test_table', 'feed_name': 'pandora_theorchard'}) assert mock_registered_executors.call_count == 0 def test_load_staging_raw_table_sme(mock_registered_executors): """Test load_staging_raw_table task.""" _run_task(tasks.load_staging_raw_table, { 'date': '2017-01-01', 'processed_datetime': '2017-10-01T00:00:00', 'filename': 'testfilename', 'temp_streams_table': 'streams_test_table', 'temp_metadata_table': 'metadata_test_table', 'feed_name': 'pandora_sme'}) (mock_registered_executors.load_staging_raw_table.assert_called_with( '2017-01-01', '2017-10-01T00:00:00', 'testfilename', config.snowflake_table_names['staging_raw'], 'streams_test_table', 'metadata_test_table', date_for_sqlloader='2017-01-01')) mock_registered_executors.reset_mock() _run_completed_task(tasks.load_staging_raw_table, { 'date': '2017-01-01', 'processed_datetime': '2017-10-01T00:00:00', 'filename': 'testfilename', 'temp_streams_table': 'streams_test_table', 'temp_metadata_table': 'metadata_test_table', 'feed_name': 'pandora_sme'}) assert mock_registered_executors.call_count == 0 def test_drop_temp_table(mock_executor_context): """Test drop_temp_table task.""" _run_task(tasks.drop_temp_table, {'temp_table_name': 'test_table'}) (mock_executor_context.drop_table.assert_called_with( 'test_table'))