def test_execute_with_gzip(self, gcs_mock_hook, azure_fileshare_mock_hook):
"""Test the execute function when the run is successful."""
operator = AzureFileShareToGCSOperator(
task_id=TASK_ID,
share_name=AZURE_FILESHARE_SHARE,
directory_name=AZURE_FILESHARE_DIRECTORY_NAME,
wasb_conn_id=WASB_CONN_ID,
gcp_conn_id=GCS_CONN_ID,
dest_gcs=GCS_PATH_PREFIX,
google_impersonation_chain=IMPERSONATION_CHAIN,
gzip=True,
)
azure_fileshare_mock_hook.return_value.list_files.return_value = MOCK_FILES
operator.execute(None)
gcs_mock_hook.return_value.upload.assert_has_calls(
[
mock.call('gcs-bucket', 'data/TEST1.csv', mock.ANY, gzip=True),
mock.call('gcs-bucket', 'data/TEST3.csv', mock.ANY, gzip=True),
mock.call('gcs-bucket', 'data/TEST2.csv', mock.ANY, gzip=True),
],
any_order=True,
)
def test_init(self):
"""Test AzureFileShareToGCSOperator instance is properly initialized."""
operator = AzureFileShareToGCSOperator(
task_id=TASK_ID,
share_name=AZURE_FILESHARE_SHARE,
directory_name=AZURE_FILESHARE_DIRECTORY_NAME,
wasb_conn_id=WASB_CONN_ID,
gcp_conn_id=GCS_CONN_ID,
dest_gcs=GCS_PATH_PREFIX,
google_impersonation_chain=IMPERSONATION_CHAIN,
)
assert operator.task_id == TASK_ID
assert operator.share_name == AZURE_FILESHARE_SHARE
assert operator.directory_name == AZURE_FILESHARE_DIRECTORY_NAME
assert operator.wasb_conn_id == WASB_CONN_ID
assert operator.gcp_conn_id == GCS_CONN_ID
assert operator.dest_gcs == GCS_PATH_PREFIX
assert operator.google_impersonation_chain == IMPERSONATION_CHAIN
AZURE_DIRECTORY_NAME = "test-azure-dir"
with DAG(
dag_id='azure_fileshare_to_gcs_example',
default_args={
'owner': 'airflow',
'depends_on_past': False,
'email': ['*****@*****.**'],
'email_on_failure': False,
'email_on_retry': False,
'retries': 1,
'retry_delay': timedelta(minutes=5),
},
schedule_interval=None,
start_date=datetime(2018, 11, 1),
tags=['example'],
) as dag:
# [START howto_operator_azure_fileshare_to_gcs_basic]
sync_azure_files_with_gcs = AzureFileShareToGCSOperator(
task_id='sync_azure_files_with_gcs',
share_name=AZURE_SHARE_NAME,
dest_gcs=DEST_GCS_BUCKET,
directory_name=AZURE_DIRECTORY_NAME,
azure_fileshare_conn_id='azure_fileshare_default',
gcp_conn_id='google_cloud_default',
replace=False,
gzip=True,
google_impersonation_chain=None,
)
# [END howto_operator_azure_fileshare_to_gcs_basic]