def test_execute(self, gcs_mock_hook, adls_one_mock_hook, adls_two_mock_hook): """Test the execute function when the run is successful.""" operator = AdlsToGoogleCloudStorageOperator( task_id=TASK_ID, src_adls=ADLS_PATH_1, dest_gcs=GCS_PATH, replace=False, azure_data_lake_conn_id=AZURE_CONN_ID, google_cloud_storage_conn_id=GCS_CONN_ID ) adls_one_mock_hook.return_value.list.return_value = MOCK_FILES adls_two_mock_hook.return_value.list.return_value = MOCK_FILES def _assert_upload(bucket, object, filename): gcs_bucket, gcs_object_path = _parse_gcs_url(GCS_PATH) self.assertEqual(gcs_bucket, 'test') self.assertIn(object[len(gcs_object_path):], MOCK_FILES) gcs_mock_hook.return_value.upload.side_effect = _assert_upload uploaded_files = operator.execute(None) adls_one_mock_hook.assert_called_once_with(azure_data_lake_conn_id=AZURE_CONN_ID) adls_two_mock_hook.assert_called_once_with(azure_data_lake_conn_id=AZURE_CONN_ID) gcs_mock_hook.assert_called_once_with( google_cloud_storage_conn_id=GCS_CONN_ID, delegate_to=None) # we expect MOCK_FILES to be uploaded self.assertEqual(sorted(MOCK_FILES), sorted(uploaded_files))
def test_execute(self, gcs_mock_hook, adls_one_mock_hook, adls_two_mock_hook): """Test the execute function when the run is successful.""" operator = AdlsToGoogleCloudStorageOperator( task_id=TASK_ID, src_adls=ADLS_PATH_1, dest_gcs=GCS_PATH, replace=False, azure_data_lake_conn_id=AZURE_CONN_ID, google_cloud_storage_conn_id=GCS_CONN_ID) adls_one_mock_hook.return_value.list.return_value = MOCK_FILES adls_two_mock_hook.return_value.list.return_value = MOCK_FILES def _assert_upload(bucket_name, object_name, filename): # pylint: disable=unused-argument gcs_bucket, gcs_object_path = _parse_gcs_url(GCS_PATH) self.assertEqual(gcs_bucket, 'test') self.assertIn(object_name[len(gcs_object_path):], MOCK_FILES) gcs_mock_hook.return_value.upload.side_effect = _assert_upload uploaded_files = operator.execute(None) adls_one_mock_hook.assert_called_once_with( azure_data_lake_conn_id=AZURE_CONN_ID) adls_two_mock_hook.assert_called_once_with( azure_data_lake_conn_id=AZURE_CONN_ID) gcs_mock_hook.assert_called_once_with( google_cloud_storage_conn_id=GCS_CONN_ID, delegate_to=None) # we expect MOCK_FILES to be uploaded self.assertEqual(sorted(MOCK_FILES), sorted(uploaded_files))
def test_execute(self, gcs_mock_hook, adls_one_mock_hook, adls_two_mock_hook): """Test the execute function when the run is successful.""" operator = AdlsToGoogleCloudStorageOperator( task_id=TASK_ID, src_adls=ADLS_PATH_1, dest_gcs=GCS_PATH, replace=False, azure_data_lake_conn_id=AZURE_CONN_ID, google_cloud_storage_conn_id=GCS_CONN_ID) adls_one_mock_hook.return_value.list.return_value = MOCK_FILES adls_two_mock_hook.return_value.list.return_value = MOCK_FILES # gcs_mock_hook.return_value.upload.side_effect = _assert_upload uploaded_files = operator.execute(None) gcs_mock_hook.return_value.upload.assert_has_calls([ mock.call(bucket_name='test', filename=mock.ANY, object_name='test/path/PARQUET.parquet', gzip=False), mock.call(bucket_name='test', filename=mock.ANY, object_name='test/path/TEST3.csv', gzip=False), mock.call(bucket_name='test', filename=mock.ANY, object_name='test/path/PIC.png', gzip=False), mock.call(bucket_name='test', filename=mock.ANY, object_name='test/TEST1.csv', gzip=False), mock.call(bucket_name='test', filename=mock.ANY, object_name='test/TEST2.csv', gzip=False) ], any_order=True) adls_one_mock_hook.assert_called_once_with( azure_data_lake_conn_id=AZURE_CONN_ID) adls_two_mock_hook.assert_called_once_with( azure_data_lake_conn_id=AZURE_CONN_ID) gcs_mock_hook.assert_called_once_with( google_cloud_storage_conn_id=GCS_CONN_ID, delegate_to=None) # we expect MOCK_FILES to be uploaded self.assertEqual(sorted(MOCK_FILES), sorted(uploaded_files))