"""Unit tests for the tasks of the Cable Ingestion ETL.""" from datetime import datetime from os import path from random import shuffle from tempfile import NamedTemporaryFile from tempfile import TemporaryDirectory from unittest.mock import Mock from unittest.mock import patch from flows.cable_ingestion import status from flows.cable_ingestion import tasks from flows.exceptions import EmptyGeneratorError @patch('flows.cable_ingestion.tasks.etl_util') @patch('builtins.open') @patch('flows.cable_ingestion.tasks.pysftp') @patch('flows.cable_ingestion.tasks.log') @patch('flows.cable_ingestion.tasks.os') @patch('flows.cable_ingestion.tasks.shutil') @patch('flows.cable_ingestion.tasks.tempfile') def test_download_from_server( tempfile, shutil, os, log, pysftp, mock_open, etl_util): """Test download_from_server function.""" context = { 'activity': Mock(), 'correlation_id': 'some-test', 'sftp_credentials': {'host': 'test'}, 'date_end': '20170101', 'date_start': '20160101'} sys_temp_dir = '/some/test/dir' tempfile.gettempdir.return_value = sys_temp_dir os.path.isdir.return_value = True temp_directory = path.join(sys_temp_dir, context['correlation_id']) valid_files = [ '20160530_20160807_20160809152932.ctl', '20160530_20160807_20160809152932.tar.gz', '20160606_20160814_20160816082920.ctl', '20160606_20160814_20160816082920.tar.gz'] all_files = valid_files[:] all_files.extend([ '20160523_20170731_20160807 164052.ctl', '20160530_20170807_20160814 102139.ctl']) filtered_file = [ { 'ctl': '20160530_20160807_20160809152932.ctl', 'tar.gz': '20160530_20160807_20160809152932.tar.gz'}, { 'ctl': '20160606_20160814_20160816082920.ctl', 'tar.gz': '20160606_20160814_20160816082920.tar.gz'}] filtered_file_list = [] for pair in filtered_file: for filename in pair.values(): filtered_file_list.append(filename) expected_args = { (filename, path.join(temp_directory, filename)) for filename in filtered_file_list} client = Mock() client.__enter__ = Mock(return_value=client) client.__exit__ = Mock(return_value=None) pysftp.Connection.return_value = client client.listdir.return_value = all_files etl_util.filter_files_by_date_range.return_value = valid_files etl_util.filter_file_pairs.return_value = filtered_file result_context = tasks.download_from_server(**context) assert result_context == {'directory': temp_directory} shutil.rmtree.assert_called_with(temp_directory) client.cwd.assert_called_with('OUT') client.listdir.assert_called_once client.get.assert_called result_calls = { call[0] for call in client.get.call_args_list} assert result_calls == expected_args log.update_status.assert_called_with( 'some-test', status.DOWNLOADED_FROM_FTP) @patch('flows.cable_ingestion.tasks.log') def test_set_final_status(log): """Testing set_final_status to testing.""" params = { 'activity': Mock(), 'correlation_id': 'cid1234'} tasks.set_final_status(**params) log.update_status.assert_called_with('cid1234', status.COMPLETED, True) @patch('builtins.open') @patch('flows.cable_ingestion.tasks.etl_util') @patch('flows.cable_ingestion.tasks.log') @patch('flows.cable_ingestion.tasks.util') def test_validate_source_files_valid(util, log, etl_util, mock_open): """Test validate_source_files execute function. This test bypasses filesystem lookups and asserts results based on ctl and rentrak metadata content. """ etl_util.filter_file_pairs.return_value = [ {'ctl': 'foo.ctl', 'tar.gz': 'foo.tar.gz'}] etl_util.extract_ctl_data.return_value = [ {'name': 'test_0.csv', 'md5': 'ab21', 'size': 9021}, {'name': 'test_1.csv', 'md5': 'ab22', 'size': 9026}, {'name': 'test_2.csv', 'md5': 'feb0', 'size': 1569}, {'name': 'test_3.csv', 'md5': '9ee5', 'size': 7794}] etl_util.extract_rentrak_metadata.return_value = [ {'name': 'test_0.csv', 'md5': 'ab21', 'size': 9021}, {'name': 'test_1.csv', 'md5': 'ab22', 'size': 9026}, {'name': 'test_2.csv', 'md5': 'feb0', 'size': 1569}, {'name': 'test_3.csv', 'md5': '9ee5', 'size': 7794}] # using set due to unsure dict key ordering expected = {'valid_files': {'foo.tar.gz', 'foo.ctl'}} results = tasks.validate_source_files(None, 'correlation-id-1234', '') results['valid_files'] = set(results['valid_files']) assert expected == results log.update_status.assert_called_with( 'correlation-id-1234', status.VALIDATED_SOURCE_FILES) assert not util.send_sns_message.called @patch('builtins.open') @patch('flows.cable_ingestion.tasks.etl_util') @patch('flows.cable_ingestion.tasks.log') @patch('flows.cable_ingestion.tasks.util') def test_validate_source_files_invalid(util, log, etl_util, mock_open): """Test validate_source_files execute function. This test bypasses filesystem lookups and asserts results based on ctl and rentrak metadata content. Here, even one tiny difference considers the pair as invalid. """ etl_util.filter_file_pairs.return_value = [ {'ctl': 'foo.ctl', 'tar.gz': 'foo.tar.gz'}] etl_util.extract_ctl_data.return_value = [ {'name': 'test_0.csv', 'md5': 'ab21', 'size': 9021}, {'name': 'test_1.csv', 'md5': 'ab22', 'size': 9026}, {'name': 'test_2.csv', 'md5': 'feb0', 'size': 1569}, {'name': 'test_3.csv', 'md5': '9ee5', 'size': 7794}] etl_util.extract_rentrak_metadata.return_value = [ {'name': 'test_0.csv', 'md5': 'ab21', 'size': 9022}, {'name': 'test_1.csv', 'md5': 'ab22', 'size': 9026}, {'name': 'test_2.csv', 'md5': 'feb0', 'size': 1569}, {'name': 'test_3.csv', 'md5': '9ee5', 'size': 7794}] expected = {'valid_files': False} results = tasks.validate_source_files(None, 'correlation-id-1234', '') assert expected == results log.update_status.assert_not_called() assert util.send_sns_message.called @patch('builtins.open') @patch('flows.cable_ingestion.tasks.log') @patch('flows.cable_ingestion.tasks.s3') def test_upload_to_s3_archive(s3, log, mock_open): """Test upload_to_s3_archive function.""" bucket = 'test-bucket' correlation_id = '0123-4567-8910-1112' destination = Mock() destination.format.return_value = '/some/location/' valid_files = ['file.1', 'file.2', 'file.3', 'file.4', 'file.5'] s3_objects = [Mock() for i in valid_files] s3.get_object.side_effect = s3_objects tasks.upload_to_s3_archive( None, bucket, correlation_id, destination, valid_files) destination.format.assert_called_with( bucket=bucket, correlation_id=correlation_id) for index, filename in enumerate(valid_files): s3_call_param = path.join('/some/location/', filename) s3.get_object.assert_any_call(s3_call_param) mock_open.assert_any_call(filename, 'rb') assert s3_objects[index].put.called log.update_status.assert_called_with( correlation_id, status.UPLOADED_RAW_TO_S3_ARCHIVES) @patch('builtins.open') @patch('flows.cable_ingestion.tasks.log') @patch('flows.cable_ingestion.tasks.s3') def test_upload_to_s3_archive_skip(s3, log, mock_open): """Test skip due to no valid files from previous task.""" bucket = 'test-bucket' correlation_id = '0123-4567-8910-1112' destination = Mock() destination.format.return_value = '/some/location/' valid_files = [] tasks.upload_to_s3_archive( None, bucket, correlation_id, destination, valid_files) destination.format.assert_not_called() log.update_status.assert_not_called() @patch('flows.cable_ingestion.tasks.log') def test_cleanup_local_raw(log): """Test cleanup_local_raw function.""" correlation_id = '2111-0198-7654-3210' try: with TemporaryDirectory() as td: some_file = NamedTemporaryFile(delete=False, dir=td) tasks.cleanup_local_raw( None, correlation_id, td, {'valid_files': '123'}) assert not path.isfile(some_file.name) try: path.exists(td) except FileNotFoundError: pass except: assert False except FileNotFoundError: # wrapped one more time due to another exception raised from the # TemporaryDirectory() context exiting with a missing directory. pass log.update_status.assert_called_with( correlation_id, status.CLEANED_UP_LOCAL_RAW) @patch('flows.cable_ingestion.tasks.log') def test_cleanup_local_raw_skip(log): """Test skip due to no valid files from previous task.""" correlation_id = '2111-0198-7654-3210' tasks.cleanup_local_raw( None, correlation_id, '/it/does/not/matter/', False) assert not log.update_status.called @patch('flows.cable_ingestion.tasks.datastore') @patch('flows.cable_ingestion.tasks.log') def test_create_temp_raw_table(log, datastore): """Test create_temp_raw_table function.""" correlation_id = '0123-4567-8910-1112' sql = 'TEST QUERY FOR {table_name}' temp_table_name = 'some_test_table_{correlation_hex}' tasks.create_temp_raw_table(None, correlation_id, sql, temp_table_name) datastore.execute.assert_called_with( 'TEST QUERY FOR some_test_table_0123456789101112') log.update_status.assert_called_with( correlation_id, status.CREATED_TEMP_RAW_TABLE) @patch('flows.cable_ingestion.tasks.datastore') @patch('flows.cable_ingestion.tasks.etl_util') @patch('flows.cable_ingestion.tasks.log') def test_insert_to_temp_raw_table(log, etl_util, datastore, database_context): """Test insert_to_temp_raw_table function.""" correlation_id = '0123-4567-8910-1112' insert_query = Mock() insert_sql = Mock() insert_sql.format.return_value = insert_query delete_query = Mock() delete_sql = Mock() delete_sql.format.return_value = delete_query s3_files_sorted = [str(i) + '.tar.gz' for i in range(4)] s3_files_dirty = s3_files_sorted[:] + ['foo.ctl', 'bar.ext', 'baz.fake'] shuffle(s3_files_dirty) table_name = None tarball_rows = [ [123, 234, 345, 456], ['qwe', 'wer', 'ert', 'rty'], ['asd', 'sdf', 'dfg', 'fgh'], ['zxc', 'xcv', 'cvb', 'vbn']] etl_util.get_rows_from_tarball.side_effect = tarball_rows datastore.context = database_context filename_dates = {'date_start': 'foo', 'date_end': 'baz'} etl_util.extract_filename_dates.return_value = filename_dates tasks.insert_to_temp_raw_table( None, correlation_id, delete_sql, insert_sql, s3_files_dirty, table_name) database_context._cursor.execute.assert_any_call( delete_query, filename_dates) for row in tarball_rows: database_context._cursor.executemany.assert_any_call(insert_query, row) # we shuffled s3_files but get_rows should be called sorted util_calls = enumerate(etl_util.get_rows_from_tarball.call_args_list) for i, util_call in util_calls: assert util_call[0][0] == s3_files_sorted[i] log.update_status.assert_called_with( correlation_id, status.INSERTED_TO_TEMP_RAW_TABLE) @patch('flows.cable_ingestion.tasks.datastore') @patch('flows.cable_ingestion.tasks.etl_util') @patch('flows.cable_ingestion.tasks.log') def test_insert_to_temp_raw_table_partial( log, etl_util, datastore, database_context): """Test insert_to_temp_raw_table function.""" correlation_id = '0123-4567-8910-1112' insert_query = Mock() insert_sql = Mock() insert_sql.format.return_value = insert_query delete_query = Mock() delete_sql = Mock() delete_sql.format.return_value = delete_query s3_files = [str(i) + '.tar.gz' for i in range(5)] table_name = None tarball_rows = [ [123, 234, 345, 456], EmptyGeneratorError, ['asd', 'sdf', 'dfg', 'fgh'], EmptyGeneratorError, ['zxc', 'xcv', 'cvb', 'vbn']] etl_util.get_rows_from_tarball.side_effect = tarball_rows datastore.context = database_context filename_dates = {'date_start': 'foo', 'date_end': 'baz'} etl_util.extract_filename_dates.return_value = filename_dates tasks.insert_to_temp_raw_table( None, correlation_id, delete_sql, insert_sql, s3_files, table_name) database_context._cursor.execute.assert_any_call( delete_query, filename_dates) for row in tarball_rows: if row is EmptyGeneratorError: continue database_context._cursor.executemany.assert_any_call(insert_query, row) assert database_context._cursor.executemany.call_count == 3 log.update_status.assert_called_with( correlation_id, status.INSERTED_TO_TEMP_RAW_TABLE) @patch('flows.cable_ingestion.tasks.datastore') @patch('flows.cable_ingestion.tasks.etl_util') @patch('flows.cable_ingestion.tasks.log') def test_insert_select_to_raw_table( log, etl_util, datastore, database_context): """Test insert_to_temp_raw_table function.""" correlation_id = '0123-4567-8910-1112' date_end = '2016-01-01' date_start = '2015-01-01' raw_table_name = 'test_raw_table' temp_table_name = 'test_temp_table' date_query = 'get date range for this {table_name}' delete = 'delete some data before inserting {table_name}' drop = 'drop this table over here {table_name}' insert = 'insert data into this table {table_name}' date_query_formatted = 'get date range for this test_temp_table' delete_query = 'delete some data before inserting test_raw_table' drop_query = 'drop this table over here test_temp_table' insert_query = 'insert data into this table test_temp_table' date_query_result = Mock() date_query_result.fetchone.return_value = (date_start, date_end) datastore.query.return_value = date_query_result datastore.context = database_context tasks.insert_select_to_raw_table( None, correlation_id, delete, drop, insert, raw_table_name, date_query, temp_table_name) datastore.query.assert_any_call(date_query_formatted) database_context._cursor.execute.assert_any_call( delete_query, {'date_end': date_end, 'date_start': date_start}) database_context._cursor.execute.assert_any_call(insert_query) database_context._cursor.execute.assert_any_call(drop_query) log.update_status.assert_called_with( correlation_id, status.INSERTED_TO_RAW_TABLE) @patch('flows.cable_ingestion.tasks.config') @patch('flows.cable_ingestion.tasks.garcon_feed_status') @patch('flows.cable_ingestion.tasks.log') def test_update_dashboard_status(log, garcon_feed_status, config): """Test update_dashboard_status function.""" correlation_id = '0123-4567-8910-1112' status = 'testing' config.SWF_WORKFLOW_NAME = 'test-workflow' log.get_etl_log.return_value = [ 'correlation_id value', 'workflow_run_id value', 'date_start value', 'date_end value', datetime(2000, 1, 1), 'etl_end value', 'etl_status value'] tasks.update_dashboard_status(None, correlation_id, status) assert garcon_feed_status.set_overall_status.called_once_with( 'test-workflow', '2000-01-01', status)