"""Unit tests for tasks.""" from unittest.mock import Mock from unittest.mock import patch from pytest import fixture from flows.cable_calculation import status from flows.cable_calculation import tasks @fixture def upcs(): """Fixture for UPCs used in unload task tests.""" return ['123', '234', '345', '456', '567'] @fixture def est_dates(): """Fixture for est dates used in unload task tests.""" return { '123': '2000-01-01', '234': '2000-01-02', '345': '2000-01-03', '456': '2000-01-04', '567': '2000-01-05'} @fixture def est_dates_batches(): """Fixture for batched est dates used in unload task tests.""" return [ {'123': '2000-01-01', '234': '2000-01-02', '345': '2000-01-03'}, {'456': '2000-01-04', '567': '2000-01-05'}] @fixture def dbo(): """Fixture for dbo rows in unload task tests.""" return { '123': '1.23', '234': '2.34', '345': '3.45', '456': '4.56', '567': '5.67'} @fixture def dbo_batches(): """Fixture for batched dbo rows used in unload task tests.""" return [ {'123': '1.23', '234': '2.34', '345': '3.45'}, {'456': '4.56', '567': '5.67'}] @fixture def raw_transactions(): """Fixture for raw transactions used in unload task tests.""" return [ ['transaction 1'], ['transaction 2'], ['transaction 3'], ['transaction 4'], ['transaction 5'], ['transaction 6'], ['transaction 7'], ['transaction 8'], ['transaction 9']] @fixture def raw_transactions_batches(): """Fixture for batched raw transactions used in unload task tests.""" return [ [ ['transaction 1'], ['transaction 2'], ['transaction 3'], ['transaction 4'], ['transaction 5']], [ ['transaction 6'], ['transaction 7'], ['transaction 8'], ['transaction 9']]] @patch('flows.cable_calculation.tasks.datastore') @patch('flows.cable_calculation.tasks.log') def test_create_temp_table(log, datastore): """Test creation of temp table.""" # Mock setup create_table_sql = 'CREATE TABLE test ...' create = Mock() create.format.return_value = create_table_sql temp_table_name = Mock() temp_table_name.format.return_value = 'test_table_cid1234' params = { 'activity': Mock(), 'correlation_id': 'cid-1234.1', 'create': create, 'temp_table_name': temp_table_name} result = tasks.create_temp_table(**params) # Assertions temp_table_name.format.assert_called_with(correlation_hex='cid1234') assert result['table_name'] == 'test_table_cid1234' create.format.assert_called_with(table_name='test_table_cid1234') datastore.execute.assert_called_with(create_table_sql) log.update_status.assert_called_with( 'cid-1234.1', status.TEMP_TABLE_CREATED) @patch('flows.cable_calculation.tasks.datastore') @patch('flows.cable_calculation.tasks.log') @patch('flows.cable_calculation.tasks.queries') def test_move_temp_table_to_cable_revenue( queries, log, datastore, database_context): """Test moving data from temp table to cable revenue table.""" # Mock setup upcs = ['123', '234', '345', '456', '567'] temp_table_name = 'test-temp-table' drop_sql = Mock() queries.DROP_TEMP_TABLE.format.return_value = drop_sql insert_sql = Mock() queries.INSERT_FROM_TEMP_TABLE.format.return_value = insert_sql bad_rows_sql = Mock() queries.SELECT_BAD_ROWS_FROM_TEMP.format.return_value = bad_rows_sql delete_sqls_1 = Mock() delete_sqls_2 = Mock() delete_sqls_mock = [delete_sqls_1, delete_sqls_2] queries.get_delete_from_cable_revenue_sql.return_value = delete_sqls_mock correlation_id = '1234-5678-9101-1123' date_params = {'date_end': '2016-01-01', 'date_start': '2015-01-01'} datastore.context = database_context datastore.context._cursor.fetchone.return_value = [0] tasks.queries.get_delete_from_cable_revenue_sql.batch_size = 3 tasks.move_temp_table_to_cable_revenue( activity=Mock(), correlation_id=correlation_id, temp_table_name=temp_table_name, upcs=upcs, **date_params) # Assertions executes = database_context._cursor.execute.call_args_list queries.get_delete_from_cable_revenue_sql.assert_called_with(upcs) assert executes[0][0][0] == bad_rows_sql assert executes[1][0][0] == 'START TRANSACTION' assert executes[2][0][0] == delete_sqls_1 assert executes[2][0][1] == date_params assert executes[3][0][0] == delete_sqls_2 assert executes[3][0][1] == date_params assert executes[4][0][0] == insert_sql assert executes[5][0][0] == drop_sql log.update_status.assert_called_with( correlation_id, status.TEMP_TABLE_MADE_LIVE) def test_unload_calculate_load_temp_table( monkeypatch, upcs, est_dates, est_dates_batches, dbo, dbo_batches, raw_transactions, raw_transactions_batches): """Test load_temp_table.""" correlation_id = '1234-5678-9101-1123' date_start = '2016-08-02' date_end = '2016-09-25' select_est_dates_q = Mock() select_dbo_q = Mock() insert_q = Mock() temp_table_name = 'temp_name' select_q = Mock() # Mocking log_mock = Mock() monkeypatch.setattr(tasks, 'log', log_mock) load_est_date_for_releases_mock = Mock(return_value=est_dates_batches) monkeypatch.setattr( tasks.cable_util, 'load_est_date_for_releases', load_est_date_for_releases_mock) load_dbo_from_theatrical_revenue_mock = Mock(return_value=dbo_batches) monkeypatch.setattr( tasks.cable_util, 'load_dbo_from_theatrical_revenue', load_dbo_from_theatrical_revenue_mock) unload_from_raw_table_mock = Mock(return_value=raw_transactions_batches) monkeypatch.setattr( tasks.cable_util, 'unload_from_raw_table', unload_from_raw_table_mock) calculate_split_mock = Mock(return_value='calculated_transactions') monkeypatch.setattr( tasks.cable_util, 'calculate_split', calculate_split_mock) load_temp_table_mock = Mock(return_value='loaded') monkeypatch.setattr( tasks.cable_util, 'load_temp_table', load_temp_table_mock) results = tasks.unload_calculate_load_temp_table( activity=Mock(), upcs=upcs, correlation_id=correlation_id, select_est_dates=select_est_dates_q, select_dbo=select_dbo_q, date_start=date_start, date_end=date_end, select=select_q, insert=insert_q, temp_table_name=temp_table_name) # Assertions assert results is None load_est_date_for_releases_mock.assert_called_with( upcs, select_est_dates_q) load_dbo_from_theatrical_revenue_mock.assert_called_with( upcs, select_dbo_q) unload_from_raw_table_mock.assert_called_with( upcs, date_start, date_end, select_q) calculate_split_mock.assert_called_with( est_dates, dbo, raw_transactions) load_temp_table_mock.assert_called_with( insert_q, temp_table_name, 'calculated_transactions') assert log_mock.update_status.call_count == 3 log_mock.update_status.assert_any_call( correlation_id, status.RAW_UNLOADED) log_mock.update_status.assert_any_call( correlation_id, status.CALCULATION_COMPLETED) log_mock.update_status.assert_any_call( correlation_id, status.TEMP_TABLE_INSERTED) def test_unload_calculate_load_temp_table_stopped( monkeypatch, upcs, est_dates, est_dates_batches, dbo, dbo_batches): """Test load_temp_table with stop flag.""" correlation_id = '1234-5678-9101-1123' date_start = '2016-08-02' date_end = '2016-09-25' select_est_dates_q = Mock() select_dbo_q = Mock() insert_q = Mock() temp_table_name = 'temp_name' select_q = Mock() # Mocking log_mock = Mock() monkeypatch.setattr(tasks, 'log', log_mock) load_est_date_for_releases_mock = Mock(return_value=est_dates_batches) monkeypatch.setattr( tasks.cable_util, 'load_est_date_for_releases', load_est_date_for_releases_mock) load_dbo_from_theatrical_revenue_mock = Mock(return_value=dbo_batches) monkeypatch.setattr( tasks.cable_util, 'load_dbo_from_theatrical_revenue', load_dbo_from_theatrical_revenue_mock) unload_from_raw_table_mock = Mock(return_value=[]) monkeypatch.setattr( tasks.cable_util, 'unload_from_raw_table', unload_from_raw_table_mock) calculate_split_mock = Mock() monkeypatch.setattr( tasks.cable_util, 'calculate_split', calculate_split_mock) load_temp_table_mock = Mock() monkeypatch.setattr( tasks.cable_util, 'load_temp_table', load_temp_table_mock) results = tasks.unload_calculate_load_temp_table( activity=Mock(), upcs=upcs, correlation_id=correlation_id, select_est_dates=select_est_dates_q, select_dbo=select_dbo_q, date_start=date_start, date_end=date_end, select=select_q, insert=insert_q, temp_table_name=temp_table_name) # Assertions assert results == {'stop': True} load_est_date_for_releases_mock.assert_called_with( upcs, select_est_dates_q) load_dbo_from_theatrical_revenue_mock.assert_called_with( upcs, select_dbo_q) unload_from_raw_table_mock.assert_called_with( upcs, date_start, date_end, select_q) calculate_split_mock.assert_not_called() load_temp_table_mock.assert_not_called() assert log_mock.update_status.call_count == 2 log_mock.update_status.assert_any_call(correlation_id, status.RAW_UNLOADED) log_mock.update_status.assert_any_call( correlation_id, status.CALCULATION_NOT_NEEDED) @patch('flows.cable_calculation.tasks.log') def test_set_final_status(log): """Testing set_final_status.""" params = { 'activity': Mock(), 'correlation_id': 'cid1234'} tasks.set_final_status(**params) log.update_status.assert_called_with('cid1234', status.COMPLETED, True) @patch('flows.cable_calculation.tasks.cable_util') def test_send_sns(util): """Testing the send_sns task.""" activity = Mock() correlation_id = '1234-5678-9101-1120.1' sns_correlation_id = '1234-5678-9101-1120.1.1' date_end = 'foo' date_start = 'bar' upcs = ['123', '456'] tasks.send_sns(activity, correlation_id, date_end, date_start, upcs) util.send_success_notification.assert_called_with( sns_correlation_id, date_end, date_start, upcs) @patch('flows.cable_calculation.tasks.cable_util') def test_queue_build_cache(util): """Testing the queue_build_cache task.""" activity = Mock() correlation_id = '1234-5678-9101-1120.1' sqs_correlation_id = '1234-5678-9101-1120.1' date_end = 'foo' date_start = 'bar' upcs = ['123', '456'] tasks.queue_build_cache( activity, correlation_id, date_end, date_start, upcs) util.queue_build_cache.assert_called_with( sqs_correlation_id, date_end, date_start, upcs) @patch('flows.cable_calculation.tasks.garcon_feed_status') def test_update_dashboard_status(garcon_feed_status): """Test update_dashboard_status task.""" tasks.update_dashboard_status( activity=Mock(), wflow_name='test-workflow', end_date='2000-01-02', status='test-status') assert garcon_feed_status.set_overall_status.called_once_with( 'test-workflow', '2000-01-01', 'test-status')