"""Unit tests for Sales data tasks.""" import json from unittest.mock import Mock from unittest.mock import patch from botocore.exceptions import ClientError import pytest from flows.sales_data import config from flows.sales_data import status from flows.sales_data import tasks @pytest.fixture() def batch_upc_queries(): """Mock queries returned from any batch_param_calls util call.""" yield ['query for batch 1', 'query batch 2'] @patch('flows.sales_data.tasks.ows_accounting') @patch('flows.sales_data.tasks.log') @patch('flows.sales_data.tasks.util') def test_bootstrap_accounting_period_default( mock_util, mock_log, mock_accounting): """Test bootstrap with default lookup accounting_period_id.""" period_id = 123 mock_accounting.get_accounting_period.return_value = period_id period_data = (period_id, '2017', '02', '1') mock_util.get_period_data.return_value = period_data mock_util.has_ingested_accounting_period_id.return_value = False actual = tasks.bootstrap(Mock(), 'cc123', None) assert actual == { 'date_start': '2017-02-01', 'date_end': '2017-02-28', 'accounting_period_id': period_id} mock_log.update_status.assert_not_called() mock_log.accounting_period_update.assert_called_once_with( 'cc123', period_id) @patch('flows.sales_data.tasks.ows_accounting') @patch('flows.sales_data.tasks.log') @patch('flows.sales_data.tasks.util') def test_bootstrap_accounting_period_default_processed( mock_util, mock_log, mock_accounting): """Test bootstrap with lookup accounting_period_id already processed.""" period_id = 123 mock_accounting.get_accounting_period.return_value = period_id period_data = (period_id, '2017', '02', '1') mock_util.get_period_data.return_value = period_data mock_util.has_ingested_accounting_period_id.return_value = True actual = tasks.bootstrap(Mock(), 'cc123', None) assert actual['stop'] is True mock_log.update_status.assert_not_called() mock_log.accounting_period_update.assert_not_called() @patch('flows.sales_data.tasks.log') @patch('flows.sales_data.tasks.util') def test_bootstrap_accounting_period_manual(mock_util, mock_log): """Test bootstrap with manual valid accounting_period_id.""" period_id = 123 period_data = (period_id, '2017', '02', '1') mock_util.get_period_data.return_value = period_data actual = tasks.bootstrap(Mock(), 'cc123', period_id) assert actual == { 'date_start': '2017-02-01', 'date_end': '2017-02-28', 'accounting_period_id': period_id} mock_log.update_status.assert_not_called() @patch('flows.sales_data.tasks.log') @patch('flows.sales_data.tasks.util') def test_bootstrap_invalid_period(mock_util, mock_log): """Test bootstrap with invalid accounting_period_id.""" period_id = 500 mock_util.get_period_data.return_value = None actual = tasks.bootstrap(Mock(), 'cc123', period_id) assert actual == {'stop': True, 'reason': 'Bad accounting_period_id.'} mock_log.update_status.assert_called_once_with( 'cc123', status.TERMINATED_BAD_PERIOD_ID) @patch('flows.sales_data.tasks.datastore') def test_create_temp_table(mock_ft): """Test create_temp_table function.""" mock_ft.execute.return_value = True correlation_id = 'cc123' result = tasks.create_temp_table(Mock(), correlation_id) assert result == {'table_name': 'accounting_raw_cc123'} assert mock_ft.execute.call_count == 1 calls = mock_ft.execute.call_args_list[0][0][0] assert 'CREATE TABLE IF NOT EXISTS' in calls assert 'accounting_raw_cc123' in calls assert 'accounting_revenue_raw' in calls @patch('flows.sales_data.tasks.log') @patch('flows.sales_data.tasks.datastore') @patch('flows.sales_data.tasks.util') def test_load_raw_from_temp(mock_util, mock_ft, mock_log, database_context): """Test load_raw_from_temp function.""" mock_util.delete_raw_sql.return_value = ['delete sql'] mock_ft.context = database_context correlation_id = 'cc123' upcs = [111, 222, 333, 444] temp_table_name = 'temp_table' tasks.load_raw_from_temp( Mock(), correlation_id, upcs, temp_table_name, 200) mock_log.update_status.assert_called_with( correlation_id, status.INSERTED_RAW_TABLE) assert database_context._cursor.execute.call_count == 4 calls = database_context._cursor.execute.call_args_list assert 'START TRANSACTION' == calls[0][0][0] assert 'delete sql' == calls[1][0][0] assert 'INSERT INTO' in calls[2][0][0] assert 'accounting_revenue_raw' in calls[2][0][0] assert temp_table_name in calls[2][0][0] assert 'DROP TABLE' in calls[3][0][0] assert temp_table_name in calls[3][0][0] @patch('flows.sales_data.tasks.config') def test_prepare_sns_message(config_mock): """Test prepare_sns_message function.""" config_mock.SNS_ACTION_SUCCESS = 'test passed' config_mock.SNS_SOURCE = 'test-sales' config_mock.SNS_TOPIC_ARN = 'test-arn' correlation_id = 'cc123' accounting_period_id = 200 upcs = ['123', '456'] expected_payload = { 'action': config_mock.SNS_ACTION_SUCCESS, 'correlation_id': correlation_id, 'source': config_mock.SNS_SOURCE, 'upcs': upcs, 'accounting_period_id': accounting_period_id} result = tasks.prepare_sns_message( Mock(), correlation_id, upcs, accounting_period_id) assert result['topic'] == config_mock.SNS_TOPIC_ARN assert result['subject'] == config_mock.SNS_ACTION_SUCCESS assert json.loads(result['message']) == expected_payload @patch('flows.sales_data.tasks.datawarehouse') @patch('flows.sales_data.tasks.log') @patch('flows.sales_data.tasks.util') def test_unload_sales_data(mock_util, mock_log, mock_sf, batch_upc_queries): """Test unload_sales_data function.""" mock_util.get_unload_from_snowflake_sql.return_value = batch_upc_queries upcs = ['123', '456'] correlation_id = 'cc123' result = tasks.unload_sales_data(Mock(), correlation_id, upcs, 200) assert result == {'batch_count': 2} mock_log.update_status.assert_called_with( correlation_id, status.UNLOADED_FACT_SALES) calls = mock_sf.execute.call_args_list for each_call in calls: assert each_call[0][0] in batch_upc_queries @patch('flows.sales_data.tasks.log') @patch('flows.sales_data.tasks.datastore') @patch('flows.sales_data.tasks.flow_util') def test_load_temp_table(mock_flow_util, mock_ft, mock_log, database_context): """Test load_temp_table function.""" row_data = [['row1'], ['row2']] mock_flow_util.read_gzip_csv_from_s3.return_value = row_data mock_ft.context = database_context correlation_id = 'cc123' batch_count = len(row_data) temp_table_name = 'temp_table_{batch}' s3_path_format = config.UNLOAD_DESTINATION.format( bucket='dev-cucumbers', correlation_id=correlation_id) tasks.load_temp_table(Mock(), correlation_id, batch_count, temp_table_name) assert mock_flow_util.read_gzip_csv_from_s3.call_count == batch_count assert database_context._cursor.executemany.call_count == batch_count calls = mock_flow_util.read_gzip_csv_from_s3.call_args_list for batch_number in range(batch_count): s3_path = s3_path_format.format(batch=batch_number) assert s3_path == calls[batch_number][0][0] mock_log.update_status.assert_called_once_with( correlation_id, status.TEMP_TABLE_INSERTED) @patch('flows.sales_data.tasks.log') @patch('flows.sales_data.tasks.datastore') @patch('flows.sales_data.tasks.flow_util') def test_load_temp_table_nokey( mock_flow_util, mock_ft, mock_log, database_context): """Test load_temp_table function with one batch having no data.""" row2 = ['row2'] error = ClientError({'Error': {'Code': 'NoSuchKey'}}, 'test_operation') row_data = [error, row2] mock_flow_util.read_gzip_csv_from_s3 = Mock(side_effect=row_data) mock_ft.context = database_context correlation_id = 'cc123' batch_count = len(row_data) temp_table_name = 'temp_table_{batch}' tasks.load_temp_table(Mock(), correlation_id, batch_count, temp_table_name) assert mock_flow_util.read_gzip_csv_from_s3.call_count == batch_count assert database_context._cursor.executemany.call_count == 1 calls = database_context._cursor.executemany.call_args_list assert calls[0][0][1] == row2 mock_log.update_status.assert_called_once_with( correlation_id, status.TEMP_TABLE_INSERTED) @patch('flows.sales_data.tasks.datastore') @patch('flows.sales_data.tasks.log') @patch('flows.sales_data.tasks.util') def test_insert_daily_revenue(util, log, datastore, database_context): """Test insert_daily_revenue task.""" accounting_period_id = 111 correlation_id = '1234-3456-5678-7890' upcs = Mock() raw_rows = [Mock(), Mock()] delete_rows = [Mock(), Mock(), Mock()] insert_rows = [Mock(), Mock(), Mock()] util.get_aggregate_raw_data.return_value = raw_rows util.delete_accounting_data_sql.return_value = delete_rows util.get_daily_revenue_query_params.return_value = insert_rows datastore.context = database_context tasks.insert_daily_revenue( Mock(), correlation_id, upcs, accounting_period_id) cursor = database_context._cursor start_call = cursor.execute.call_args_list[0] delete_call = cursor.execute.call_args_list[1] insert_call = cursor.executemany.call_args_list[0] assert start_call[0][0] == 'START TRANSACTION;' assert isinstance(delete_call[0][0], Mock) assert {'accounting_period_id': 111} == delete_call[0][1] assert 'INSERT INTO' in insert_call[0][0] assert insert_call[0][1] == insert_rows log.update_status.assert_called_once_with( correlation_id, status.INSERTED_DAILY_REVENUE) @patch('flows.sales_data.tasks.log') @patch('flows.sales_data.tasks.util') def test_insert_daily_revenue_no_data(util, log): """Test insert_daily_revenue when no raw data.""" accounting_period_id = 111 correlation_id = '1234-3456-5678-7890' upcs = Mock() util.get_aggregate_raw_data.return_value = [] tasks.insert_daily_revenue( Mock(), correlation_id, upcs, accounting_period_id) log.update_status.assert_not_called() @patch('flows.sales_data.tasks.log') @patch('flows.sales_data.tasks.garcon_feed_status') def test_set_dynamo_status(feed_status_mock, log_mock): """Test set_dynamo_status task.""" correlation_id = 'cc123' date_start = '2017-01-01' tasks.set_dynamo_status(Mock(), correlation_id, date_start) log_mock.update_status.assert_called_with( correlation_id, status.DYNAMO_STATUS_UPDATED) feed_status_mock.set_overall_status.assert_called_with( config.SWF_WORKFLOW_NAME, date_start, feed_status_mock.STATUS_INGESTED) @patch('flows.sales_data.tasks.log') def test_set_final_status(log_mock): """Test to set_final_status task.""" correlation_id = 'cc123' tasks.set_final_status(Mock(), correlation_id) log_mock.update_status.assert_called_once_with( correlation_id, status.INGESTED, True)