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 = ADLSToGCSOperator(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))
def test_init(self):
"""Test AdlsToGoogleCloudStorageOperator instance is properly initialized."""
operator = ADLSToGCSOperator(task_id=TASK_ID,
src_adls=ADLS_PATH_1,
dest_gcs=GCS_PATH,
replace=False,
azure_data_lake_conn_id=AZURE_CONN_ID,
gcp_conn_id=GCS_CONN_ID)
self.assertEqual(operator.task_id, TASK_ID)
self.assertEqual(operator.src_adls, ADLS_PATH_1)
self.assertEqual(operator.dest_gcs, GCS_PATH)
self.assertEqual(operator.replace, False)
self.assertEqual(operator.gcp_conn_id, GCS_CONN_ID)
self.assertEqual(operator.azure_data_lake_conn_id, AZURE_CONN_ID)