def test_execute(self, mock_hook):
task = GCSSynchronizeBuckets(
task_id="task-id",
source_bucket="SOURCE_BUCKET",
destination_bucket="DESTINATION_BUCKET",
source_object="SOURCE_OBJECT",
destination_object="DESTINATION_OBJECT",
recursive=True,
delete_extra_files=True,
allow_overwrite=True,
gcp_conn_id="GCP_CONN_ID",
delegate_to="DELEGATE_TO",
)
task.execute({})
mock_hook.assert_called_once_with(
google_cloud_storage_conn_id='GCP_CONN_ID',
delegate_to='DELEGATE_TO')
mock_hook.return_value.sync.assert_called_once_with(
source_bucket='SOURCE_BUCKET',
source_object='SOURCE_OBJECT',
destination_bucket='DESTINATION_BUCKET',
destination_object='DESTINATION_OBJECT',
delete_extra_files=True,
recursive=True,
allow_overwrite=True,
)
BUCKET_1_SRC = os.environ.get("GCP_GCS_BUCKET_1_SRC", "test-gcs-sync-1-src")
BUCKET_1_DST = os.environ.get("GCP_GCS_BUCKET_1_DST", "test-gcs-sync-1-dst")
BUCKET_2_SRC = os.environ.get("GCP_GCS_BUCKET_2_SRC", "test-gcs-sync-2-src")
BUCKET_2_DST = os.environ.get("GCP_GCS_BUCKET_2_DST", "test-gcs-sync-2-dst")
BUCKET_3_SRC = os.environ.get("GCP_GCS_BUCKET_3_SRC", "test-gcs-sync-3-src")
BUCKET_3_DST = os.environ.get("GCP_GCS_BUCKET_3_DST", "test-gcs-sync-3-dst")
with models.DAG(
"example_gcs_to_gcs", default_args=default_args, schedule_interval=None, tags=['example']
) as dag:
sync_full_bucket = GCSSynchronizeBuckets(
task_id="sync-full-bucket",
source_bucket=BUCKET_1_SRC,
destination_bucket=BUCKET_1_DST
)
sync_to_subdirectory_and_delete_extra_files = GCSSynchronizeBuckets(
task_id="sync_to_subdirectory_and_delete_extra_files",
source_bucket=BUCKET_1_SRC,
destination_bucket=BUCKET_1_DST,
destination_object="subdir/",
delete_extra_files=True,
)
sync_from_subdirectory_and_allow_overwrite_and_non_recursive = GCSSynchronizeBuckets(
task_id="sync_from_subdirectory_and_allow_overwrite_and_non_recursive",
source_bucket=BUCKET_1_SRC,
source_object="subdir/",
BUCKET_2_SRC = os.environ.get("GCP_GCS_BUCKET_2_SRC", "test-gcs-sync-2-src")
BUCKET_2_DST = os.environ.get("GCP_GCS_BUCKET_2_DST", "test-gcs-sync-2-dst")
BUCKET_3_SRC = os.environ.get("GCP_GCS_BUCKET_3_SRC", "test-gcs-sync-3-src")
BUCKET_3_DST = os.environ.get("GCP_GCS_BUCKET_3_DST", "test-gcs-sync-3-dst")
OBJECT_1 = os.environ.get("GCP_GCS_OBJECT_1", "test-gcs-to-gcs-1")
OBJECT_2 = os.environ.get("GCP_GCS_OBJECT_2", "test-gcs-to-gcs-2")
with models.DAG("example_gcs_to_gcs",
default_args=default_args,
schedule_interval=None,
tags=['example']) as dag:
# [START howto_synch_bucket]
sync_bucket = GCSSynchronizeBuckets(task_id="sync_bucket",
source_bucket=BUCKET_1_SRC,
destination_bucket=BUCKET_1_DST)
# [END howto_synch_bucket]
# [START howto_synch_full_bucket]
sync_full_bucket = GCSSynchronizeBuckets(task_id="sync_full_bucket",
source_bucket=BUCKET_1_SRC,
destination_bucket=BUCKET_1_DST,
delete_extra_files=True,
allow_overwrite=True)
# [END howto_synch_full_bucket]
# [START howto_synch_to_subdir]
sync_to_subdirectory = GCSSynchronizeBuckets(
task_id="sync_to_subdirectory",
source_bucket=BUCKET_1_SRC,