"""Test tasks for YouTube Conflict Elasticsearch Workflow.""" from unittest.mock import MagicMock, patch import boto3 from botocore.exceptions import ClientError from garcon_contrib.dynamo_feed_status import \ garcon_feed_status as feed_status import pytest from yt_conflict_elasticsearch.flows.elasticsearch_export import config from yt_conflict_elasticsearch.flows.elasticsearch_export import tasks TASK_MODULE_PATH = 'yt_conflict_elasticsearch.flows.elasticsearch_export.tasks' GENERATORS_MODULE_PATH = ( 'yt_conflict_elasticsearch.flows.elasticsearch_export.generators') @pytest.fixture(autouse=True) def mock_os_remove(mocker): """Mock creating temp folder.""" return mocker.patch('os.remove') @pytest.fixture def mock_activity(): """Return mock activity.""" return MagicMock() @pytest.fixture() def fixture_task_date(): """Get task date fixture.""" return '2018-01-01' @pytest.fixture() def fixture_conflict_status(): """Conflict status fixture.""" return 'new' @pytest.fixture(autouse=True) def mock_s3_bucket(): """Mock s3 bucket.""" return MagicMock() @pytest.fixture(autouse=True) def mock_s3_resource(mock_s3_bucket): """Mock s3 resource that returns a mock s3 bucket.""" mock = MagicMock() mock.Bucket.return_value = mock_s3_bucket return mock @pytest.fixture(autouse=True) def mock_s3_client(mocker): """Mock s3 resource that returns a mock s3 bucket.""" client_mock = MagicMock() mock = mocker.patch.object(boto3, 'client') mock.return_value = client_mock return client_mock @pytest.fixture(autouse=True) def mock_boto_resource(mocker, mock_s3_resource, mock_s3_bucket): """Mock boto resource that returns a mock s3 resource.""" mock = mocker.patch.object(boto3, 'resource') mock.return_value = mock_s3_resource return mock @pytest.fixture(autouse=True) def mock_create_tmp_path(mocker): """Mock creating temp folder.""" return mocker.patch( TASK_MODULE_PATH + '._create_tmp_path', return_value='tmp') @pytest.fixture(autouse=True) def mock_opensearch_client(): """Mock OpenSearch client.""" return MagicMock() @pytest.fixture(autouse=True) def mock_get_opensearch_client( mocker, mock_opensearch_client): """Mock get_opensearch_client.""" client_path = TASK_MODULE_PATH + '._get_opensearch_client' return mocker.patch( client_path, return_value=mock_opensearch_client) @pytest.fixture(autouse=True) def mock_task_status(mocker): """Mock completed task.""" mock_task_status = mocker.patch( 'yt_conflict_elasticsearch.' 'flows.elasticsearch_export.tasks.task_status' ) return mock_task_status @pytest.fixture() def mock_indices_settings(): """Test indices settings.""" return { 'conflicts_2017_12_29': { 'settings': { 'index': { 'creation_date': '1514552399000', 'number_of_shards': '5', 'number_of_replicas': '1', 'version': {'created': '6020399'}, 'provided_name': 'conflicts_2017_12_29', 'uuid': '-5tlpWO4S3u19FO_5lHXbQ', } } }, 'conflicts_2017_12_30': { 'settings': { 'index': { 'creation_date': '1514638799000', 'number_of_shards': '5', 'number_of_replicas': '1', 'version': {'created': '6020399'}, 'provided_name': 'conflicts_2017_12_30', 'uuid': 'f3dC378c39F1E5beaf2f', } } }, 'conflicts_2017_12_31': { 'settings': { 'index': { 'creation_date': '1514725151000', 'number_of_shards': '5', 'number_of_replicas': '1', 'version': {'created': '6020399'}, 'provided_name': 'conflicts_2017_12_31', 'uuid': 'tH8eoYzTQ7WSyuWs9zsGLQ', } } }, 'conflicts_2018_01_01': { 'settings': { 'index': { 'creation_date': '1514811599000', 'number_of_shards': '5', 'number_of_replicas': '1', 'version': {'created': '6020399'}, 'provided_name': 'conflicts_2018_01_01', 'uuid': 'e6Cf2F24811c11E8ADC0Fa', } } } } @pytest.fixture() def mock_list_objects_v2(): """Test list objects v2.""" return [ {'Key': 'folder/2018_01_01/file_new.json.gz'}, {'Key': 'folder/2018_01_01/file_new.csv.gz'}, ] def run_completed_task(task, context, mock_task_status, task_name): """Check that task runs correctly when status is set to completed.""" return_value = task(**context) mock_task_status.is_completed_task.assert_called_with( config.SWF_FLOW_NAME, context['date'], task_name) assert not mock_task_status.mark_completed_task.called return return_value def run_not_completed_task(task, context, mock_task_status, task_name): """Check that task runs correctly when status is set to not completed.""" return_value = task(**context) mock_task_status.is_completed_task.assert_called_with( config.SWF_FLOW_NAME, context['date'], task_name) mock_task_status.mark_completed_task.assert_called_with( config.SWF_FLOW_NAME, context['date'], task_name) return return_value class TestTaskBootstrap(object): """Test for bootstrap task.""" @pytest.fixture(autouse=True) def run_task(self, mock_activity, fixture_task_date): """Run the task with mocks. This is auto runned for each test.""" return tasks.bootstrap(mock_activity, fixture_task_date, reload=False) @pytest.fixture(autouse=True) def run_task_reload(self, mock_activity, fixture_task_date): """Run the task with mocks. This is auto runned for each test.""" mock_activity.patch(feed_status, 'delete_status') return tasks.bootstrap(mock_activity, fixture_task_date, reload=True) def test_bootstrap(self, run_task, run_task_reload, mock_activity, fixture_task_date): """Test bootstrap task.""" assert run_task['date'] == fixture_task_date assert run_task_reload['date'] == fixture_task_date assert mock_activity.hasBeenCalledOnce() class TestTaskStoreCsvToS3Conflicts(object): """Tests for task to store conflict data from snowflake to CSV file.""" @pytest.fixture(autouse=True) def run_task( self, mock_activity, fixture_task_date, mock_sf_executor_context, fixture_conflict_status, mock_task_status): """Run the task with mocks. This is auto runned for each test.""" mock_sf_executor_context.store_csv_to_s3_conflicts \ .return_value = MagicMock() context = { 'activity': mock_activity, 'date': fixture_task_date, 'conflict_status': fixture_conflict_status } mock_task_status.is_completed_task.return_value = True run_completed_task( tasks.task_store_conflicts_in_database_as_csv, context, mock_task_status, 'store_conflicts_in_database_as_csv_{}'.format(fixture_conflict_status)) # noqa: E501 mock_task_status.is_completed_task.return_value = False run_not_completed_task( tasks.task_store_conflicts_in_database_as_csv, context, mock_task_status, 'store_conflicts_in_database_as_csv_{}'.format(fixture_conflict_status)) # noqa: E501 def test_snowflake_executer_called( self, mock_sf_executor_context): """Test snowflake executor is called.""" mock_sf_executor_context.store_csv_to_s3_conflicts. \ assert_called_once_with( 's3://bucket/folder/2018_01_01/file_new.csv.gz', 'new') class TestTaskTransformCsvToJsonConflicts(object): """Tests for task to create and populate Elasticsearch index.""" @pytest.fixture(autouse=True) def run_task( self, mock_activity, fixture_task_date, fixture_conflict_status, mock_task_status): """Run the task with mocks. This is auto runned for each test.""" context = { 'activity': mock_activity, 'date': fixture_task_date, 'conflict_status': fixture_conflict_status } mock_task_status.is_completed_task.return_value = True run_completed_task( tasks.task_transform_csv_to_json_conflicts, context, mock_task_status, 'transform_csv_to_json_conflicts_{}'.format(fixture_conflict_status)) # noqa: E501 mock_task_status.is_completed_task.return_value = False run_not_completed_task( tasks.task_transform_csv_to_json_conflicts, context, mock_task_status, 'transform_csv_to_json_conflicts_{}'.format(fixture_conflict_status)) # noqa: E501 @pytest.fixture(autouse=True) def mock_file_handler(self): """Return a MagicMock as a file handler.""" return MagicMock() @pytest.fixture(autouse=True) def mock_open(self, mocker, mock_file_handler): """Mock the file open call.""" return mocker.patch('builtins.open', return_value=mock_file_handler) @pytest.fixture(autouse=True) def mock_process_csv_file(self, mocker): """Mock the CSV procesing logic.""" return mocker.patch( TASK_MODULE_PATH + '.conflict_csv_to_json.process_csv_file') @pytest.fixture(autouse=True) def mock_gzip_file(self, mocker): """Mock creating temp folder.""" return mocker.patch( TASK_MODULE_PATH + '._gzip_file', return_value='tmp/file_new.json.gz') def test_temp_path_created(self, mock_create_tmp_path, fixture_task_date): """Test temp file path is created.""" mock_create_tmp_path.assert_called_with(fixture_task_date) def test_s3_resource_created(self, mock_boto_resource): """Test s3 resource is created.""" mock_boto_resource.assert_called_with('s3') def test_s3_bucket_object_created(self, mock_s3_resource): """Test the s3 bucket object is created.""" mock_s3_resource.Bucket.assert_called_with(config.S3_BUCKET) def test_csv_downloaded_from_s3_bucket(self, mock_s3_bucket): """Test conflict CSV file is downloaded from s3 bucket.""" mock_s3_bucket.download_file.assert_called_with( 'folder/2018_01_01/file_new.csv.gz', 'tmp/file_new.csv.gz') def test_json_file_is_opened(self, mock_open): """Test the JSON file is opened.""" mock_open.assert_called_with('tmp/file_new.json', 'w') def test_process_csv_file(self, mock_process_csv_file, mock_file_handler): """Test the processing of the CSV file.""" mock_process_csv_file.assert_called_with( 'tmp/file_new.csv.gz', mock_file_handler) def test_json_file_is_gzipped(self, mock_gzip_file): """Test that the JSON file gets gzipped.""" mock_gzip_file.assert_called_with('tmp/file_new.json') def test_temp_files_removed(self, mock_os_remove): """Test that temp files are removed.""" mock_os_remove.assert_any_call('tmp/file_new.csv.gz') mock_os_remove.assert_any_call('tmp/file_new.json.gz') mock_os_remove.assert_any_call('tmp/file_new.json') def test_json_uploaded_to_s3_bucket(self, mock_s3_bucket): """Test conflict JSON file is uploaded to s3 bucket.""" mock_s3_bucket.upload_file.assert_called_with( 'tmp/file_new.json.gz', 'folder/2018_01_01/file_new.json.gz') class TestTaskCreateElasticsearchIndex(object): """Tests for task to create Elasticsearch index.""" @pytest.fixture(autouse=True) def run_task( self, mock_activity, fixture_task_date, fixture_conflict_status, mock_task_status, mock_opensearch_client): """Run the task with mocks. This is auto runned for each test.""" context = { 'activity': mock_activity, 'date': fixture_task_date } mock_task_status.is_completed_task.return_value = True mock_opensearch_client.indices.exists_alias = lambda name: False mock_opensearch_client.indices.exists = lambda index: False run_completed_task( tasks.task_create_elasticsearch_index, context, mock_task_status, 'create_elasticsearch_index') mock_task_status.is_completed_task.return_value = False run_not_completed_task( tasks.task_create_elasticsearch_index, context, mock_task_status, 'create_elasticsearch_index') def test_elasticsearch_index_created(self, mock_opensearch_client): """Test Opensearch index is created.""" mock_opensearch_client.indices.create.assert_called_with( index=config.OS_CONFLICTS_INDEX_NAME) mock_opensearch_client.indices.put_alias.assert_called_with( name=config.OS_CONFLICTS_ALIAS_NAME, index=config.OS_CONFLICTS_INDEX_NAME) class TestTaskPopulateElasticsearchIndex(object): """Tests for task to populate Elasticsearch index.""" @pytest.fixture(autouse=True) def run_task( self, mock_activity, fixture_task_date, fixture_conflict_status, mock_task_status, mock_sf_executor_context): """Run the task with mocks. This is auto runned for each test.""" context = { 'activity': mock_activity, 'date': fixture_task_date, 'conflict_status': fixture_conflict_status } mock_task_status.is_completed_task.return_value = True run_completed_task( tasks.task_populate_elasticsearch_index, context, mock_task_status, 'populate_elasticsearch_index_{}'.format(fixture_conflict_status)) mock_task_status.is_completed_task.return_value = False run_not_completed_task( tasks.task_populate_elasticsearch_index, context, mock_task_status, 'populate_elasticsearch_index_{}'.format(fixture_conflict_status)) @pytest.fixture(autouse=True) def mock_opensearch_bulk(self, mocker): """Mock opensearch bulk helper function.""" res = iter([(True, 'somehash')]) return mocker.patch(TASK_MODULE_PATH + '.opensearch_bulk', return_value=res) @pytest.fixture(autouse=True) def mock_json_generator(self, mocker): """Mock JSON generator.""" return MagicMock() @pytest.fixture(autouse=True) def mock_get_conflict_data_to_populate_os( self, mocker, mock_json_generator): """Mock creating temp folder.""" return mocker.patch( GENERATORS_MODULE_PATH + '.get_conflict_data_to_populate_os', return_value=mock_json_generator) @pytest.fixture(autouse=True) def mock_get_unindexed_conflict_ids_by_line(self, mocker): """Mock get conflicts ids by line generator.""" return mocker.patch( GENERATORS_MODULE_PATH + '.get_unindexed_conflict_ids_by_line') @pytest.fixture(autouse=True) def mock_get_es_ids_to_update_territories_os( self, mocker, mock_json_generator): """Mock get es_ids to update territories.""" return mocker.patch( GENERATORS_MODULE_PATH + '.get_es_ids_to_update_territories_os', return_value=mock_json_generator) @pytest.fixture(autouse=True) def mock_update_partially_resolved_conflicts_in_os(self, mocker): """Mock update partially resolved conflict.""" return mocker.patch( TASK_MODULE_PATH + '._update_partially_resolved_conflicts_in_os') @pytest.fixture(autouse=True) def mock_get_batch_to_insert_to_sf_temp_os(self, mocker): """Mock get conflicts ids by line generator.""" sf_records = [ { 'conflict_id': 1, 'es_id': 'jhgvhyygvhybu' }, { 'conflict_id': 2, 'es_id': 'jhgvhyygvhybu' }, { 'conflict_id': 3, 'es_id': 'jhgvhyygvhybu' } ] return mocker.patch( GENERATORS_MODULE_PATH + '.get_batch_to_insert_to_sf_temp_os', return_value=sf_records) def test_temp_path_created(self, mock_create_tmp_path, fixture_task_date): """Test temp file path is created.""" mock_create_tmp_path.assert_called_with(fixture_task_date) def test_s3_resource_created(self, mock_boto_resource): """Test s3 resource is created.""" mock_boto_resource.assert_called_with('s3') def test_s3_bucket_object_created(self, mock_s3_resource): """Test the s3 bucket object is created.""" mock_s3_resource.Bucket.assert_called_with(config.S3_BUCKET) def test_json_downloaded_from_s3_bucket(self, mock_s3_bucket): """Test conflict JSON file is downloaded from s3 bucket.""" mock_s3_bucket.download_file.assert_called_with( 'folder/2018_01_01/file_new.json.gz', 'tmp/file_new.json.gz') def test_temp_file_removed(self, mock_os_remove): """Test the s3 bucket object is created.""" mock_os_remove.assert_called_with('tmp/file_new.json.gz') def test_elasticsearch_index_populated( self, mock_opensearch_client, mock_opensearch_bulk, mock_get_unindexed_conflict_ids_by_line, mock_get_conflict_data_to_populate_os, mock_get_es_ids_to_update_territories_os, mock_get_batch_to_insert_to_sf_temp_os, mock_update_partially_resolved_conflicts_in_os, mock_json_generator, mock_sf_executor_context): """Test OpenSearch index is populated.""" mock_get_es_ids_to_update_territories_os.assert_called_with( 'tmp/file_new.json.gz') mock_update_partially_resolved_conflicts_in_os.assert_called() mock_get_conflict_data_to_populate_os.assert_called_with( 'tmp/file_new.json.gz', 'conflict_write') mock_get_batch_to_insert_to_sf_temp_os.assert_called() mock_sf_executor_context.create_temp_conflict_to_es_id.assert_called() mock_sf_executor_context.fill_temp_table_with_es_ids.assert_called() mock_sf_executor_context.populate_es_ids_to_conflicts.assert_called() mock_opensearch_bulk.assert_called_once_with( mock_opensearch_client, mock_json_generator, chunk_size=config.CHUNK_SIZE, max_chunk_bytes=config.MAX_CHUNK_BYTES, thread_count=config.THREAD_COUNT, queue_size=config.QUEUE_SIZE) mock_get_unindexed_conflict_ids_by_line.return_value = \ iter([[123, 345, 456], [1111, 2222]]) class TestTaskMarkIndexedConflicts(object): """Tests for task to mark indexed conflicts.""" @pytest.fixture(autouse=True) def run_task( self, mock_activity, fixture_task_date, mock_task_status, mock_opensearch_client, mock_indices_settings, mock_sf_executor_context): """Run the task with mocks. This is auto runned for each test.""" context = { 'activity': mock_activity, 'date': fixture_task_date, } mock_task_status.is_completed_task.return_value = True run_completed_task( tasks.task_mark_indexed_conflicts, context, mock_task_status, 'mark_indexed_conflicts') mock_task_status.is_completed_task.return_value = False mock_opensearch_client.indices.get_settings = MagicMock( return_value=mock_indices_settings ) run_not_completed_task( tasks.task_mark_indexed_conflicts, context, mock_task_status, 'mark_indexed_conflicts') def test_snowflake_executer_called( self, run_task, mock_sf_executor_context): """Test snowflake executor is called.""" mock_sf_executor_context.mark_indexed_conflicts.\ assert_called_once_with() class TestTaskCleanUpS3(object): """Tests for task to delete conflicts files from S3 Bucket.""" @pytest.fixture(autouse=True) def run_task( self, mock_activity, fixture_task_date, mock_task_status, mock_s3_client, mock_list_objects_v2): """Run the task with mocks. This is auto runned for each test.""" context = { 'activity': mock_activity, 'date': fixture_task_date, } mock_task_status.is_completed_task.return_value = True run_completed_task( tasks.clean_up_s3, context, mock_task_status, 'clean_up_s3') mock_task_status.is_completed_task.return_value = False mock_s3_client.list_objects_v2 = MagicMock( return_value={'Contents': mock_list_objects_v2} ) run_not_completed_task( tasks.clean_up_s3, context, mock_task_status, 'clean_up_s3') def test_conflict_files_deleted( self, mock_s3_client, fixture_task_date, mock_list_objects_v2): """Test conflicts files are deleted from S3 Bucket.""" s3_tmp_path = config.S3_FOLDER_TEMPLATE.format( date=fixture_task_date.replace('-', '_')) mock_s3_client.list_objects_v2.assert_called_with( Bucket=config.S3_BUCKET, Prefix=s3_tmp_path) mock_s3_client.delete_objects.assert_called_with( Bucket=config.S3_BUCKET, Delete={'Objects': mock_list_objects_v2} ) class TestTaskRemoveRespondedConflictsFromEs(object): """Tests for task to populate Elasticsearch index.""" @pytest.fixture(autouse=True) def run_task( self, request, mock_activity, fixture_task_date, mock_task_status, mock_sf_executor_context): """Run the task with mocks. This is auto runned for each test.""" if 'noesids' in request.keywords: mock_sf_executor_context.\ get_responded_conflicts.return_value = [] else: mock_sf_executor_context\ .get_responded_conflicts.return_value = ['test'] context = { 'activity': mock_activity, 'date': fixture_task_date } mock_task_status.is_completed_task.return_value = True run_completed_task( tasks.remove_responded_conflicts_from_es, context, mock_task_status, 'remove_responded_conflicts_from_es') mock_task_status.is_completed_task.return_value = False run_not_completed_task( tasks.remove_responded_conflicts_from_es, context, mock_task_status, 'remove_responded_conflicts_from_es') @pytest.fixture(autouse=True) def mock_opensearch_bulk(self, mocker): """Mock opensearch bulk helper function.""" return mocker.patch(TASK_MODULE_PATH + '.opensearch_bulk') @pytest.fixture(autouse=True) def mock_conflicts_generator(self, mocker): """Mock JSON generator.""" return MagicMock() @pytest.fixture(autouse=True) def mock_conflicts_to_remove_os_generator( self, mocker, mock_conflicts_generator): """Mock JSON generator.""" return mocker.patch( GENERATORS_MODULE_PATH + '.conflicts_to_remove_os', return_value=mock_conflicts_generator) def test_opensearch_items_removed( self, mock_opensearch_client, mock_opensearch_bulk, mock_conflicts_generator): """Test OpenSearch index is populated.""" mock_opensearch_bulk.assert_called_once_with( mock_opensearch_client, mock_conflicts_generator, chunk_size=config.CHUNK_SIZE, max_chunk_bytes=config.MAX_CHUNK_BYTES, raise_on_error=False, thread_count=config.THREAD_COUNT, queue_size=config.QUEUE_SIZE) def test_snowflake_executer_called( self, run_task, mock_sf_executor_context): """Test snowflake executor is called.""" mock_sf_executor_context.get_responded_conflicts.\ assert_called_once_with() mock_sf_executor_context.remove_selected_es_ids.\ assert_called_once() @pytest.mark.noesids def test_snowflake_executer_not_called( self, run_task, mock_sf_executor_context): """Test snowflake executor is not called.""" mock_sf_executor_context.get_responded_conflicts.\ assert_called_once() mock_sf_executor_context.remove_selected_es_ids.\ assert_not_called() class TestTaskRemoveResolvedConflictsFromEs(object): """Tests for task to populate Elasticsearch index.""" @pytest.fixture(autouse=True) def run_task( self, request, mock_activity, fixture_task_date, mock_task_status, mock_sf_executor_context): """Run the task with mocks. This is auto runned for each test.""" if 'noesids' in request.keywords: mock_sf_executor_context.get_resolved_conflicts.return_value = [] else: mock_sf_executor_context\ .get_resolved_conflicts.return_value = ['test'] context = { 'activity': mock_activity, 'date': fixture_task_date } mock_task_status.is_completed_task.return_value = True run_completed_task( tasks.remove_resolved_conflicts_from_es, context, mock_task_status, 'remove_resolved_conflicts_from_es') mock_task_status.is_completed_task.return_value = False run_not_completed_task( tasks.remove_resolved_conflicts_from_es, context, mock_task_status, 'remove_resolved_conflicts_from_es') @pytest.fixture(autouse=True) def mock_opensearch_bulk(self, mocker): """Mock opensearch bulk helper function.""" return mocker.patch(TASK_MODULE_PATH + '.opensearch_bulk') @pytest.fixture(autouse=True) def mock_conflicts_generator(self, mocker): """Mock JSON generator.""" return MagicMock() @pytest.fixture(autouse=True) def mock_conflicts_to_remove_os_generator( self, mocker, mock_conflicts_generator): """Mock JSON generator.""" return mocker.patch( GENERATORS_MODULE_PATH + '.conflicts_to_remove_os', return_value=mock_conflicts_generator) def test_opensearch_items_removed( self, mock_opensearch_client, mock_opensearch_bulk, mock_conflicts_generator): """Test OpenSearch index is populated.""" mock_opensearch_bulk.assert_called_once_with( mock_opensearch_client, mock_conflicts_generator, chunk_size=config.CHUNK_SIZE, max_chunk_bytes=config.MAX_CHUNK_BYTES, raise_on_error=False, thread_count=config.THREAD_COUNT, queue_size=config.QUEUE_SIZE) def test_snowflake_executer_called( self, run_task, mock_sf_executor_context): """Test snowflake executor is called.""" mock_sf_executor_context.get_resolved_conflicts.\ assert_called_once_with() mock_sf_executor_context.remove_selected_es_ids.\ assert_called_once() @pytest.mark.noesids def test_snowflake_executer_not_called( self, run_task, mock_sf_executor_context): """Test snowflake executor is not called.""" mock_sf_executor_context.get_resolved_conflicts.\ assert_called_once() mock_sf_executor_context.remove_selected_es_ids.\ assert_not_called() def test_download_file_file_exists(mock_s3_bucket): """Test _download_file returns True if file exists.""" assert tasks._download_file(mock_s3_bucket, 'key', 'path') mock_s3_bucket.download_file.assert_called_with('key', 'path') def test_download_file_file_does_not_exists(mock_s3_bucket): """Test _download_file returns False if file does not exist.""" error_response = {'Error': {'Code': '404'}} mock_s3_bucket.download_file.side_effect = ClientError( error_response, 'HeadObject') assert tasks._download_file(mock_s3_bucket, 'key', 'path') is False def test_download_file_file_exception(mock_s3_bucket): """Test _download_file raise an exception in case error is not 404.""" error_response = {'Error': {'Code': '500'}} mock_s3_bucket.download_file.side_effect = ClientError( error_response, 'HeadObject') with pytest.raises(ClientError): tasks._download_file(mock_s3_bucket, 'key', 'path') @patch(TASK_MODULE_PATH + '.opensearch_bulk') def test_update_partially_resolved_conflicts_in_os( mock_opensearch_bulk, mock_opensearch_client): """Test _update_partially_resolved_conflicts_in_os.""" conflicts_data = iter([ { '_index': config.OS_CONFLICTS_INDEX_NAME, '_type': 'document', '_op_type': 'update', '_id': 'EWAvpWcBLzAzl7XpJ9Dv', 'doc': { 'territories': [ { 'conflict_id': 4971125, 'code': 'EC', 'continent_name': 'South America', 'name': 'Ecuador', }, { 'conflict_id': 5033762, 'code': 'AR', 'continent_name': 'South America', 'name': 'Argentina' } ] } }, { '_index': config.OS_CONFLICTS_INDEX_NAME, '_type': 'document', '_op_type': 'update', '_id': 'EmAvpWcBLzAzl7XpJ9Dv', 'doc': { 'territories': [ { 'conflict_id': 5765433, 'code': 'EC', 'continent_name': 'South America', 'name': 'Ecuador', } ] } } ]) tasks._update_partially_resolved_conflicts_in_os( mock_opensearch_client, iter(conflicts_data)) mock_opensearch_bulk.assert_called_with( mock_opensearch_client, conflicts_data, chunk_size=config.CHUNK_SIZE, max_chunk_bytes=config.MAX_CHUNK_BYTES, thread_count=config.THREAD_COUNT, queue_size=config.QUEUE_SIZE)