def test_execute(self, mock_hook):
task = GoogleCloudStorageSynchronizeBuckets(
task_id="task-id",
source_bucket="SOURCE_BUCKET",
destination_bucket="DESTINATION_BUCKET",
source_object="SOURCE_OBJECT",
destination_object="DESTINATION_OBJECT",
recursive="RECURSIVE",
delete_extra_files="DELETE_EXTRA_FILES",
allow_overwrite="ALLOW_OVERWRITE",
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='DELETE_EXTRA_FILES',
recursive='RECURSIVE',
allow_overwrite="ALLOW_OVERWRITE",
)
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
) as dag:
sync_full_bucket = GoogleCloudStorageSynchronizeBuckets(
task_id="sync-full-bucket",
source_bucket=BUCKET_1_SRC,
destination_bucket=BUCKET_1_DST
)
sync_to_subdirectory_and_delete_extra_files = GoogleCloudStorageSynchronizeBuckets(
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 = GoogleCloudStorageSynchronizeBuckets(
task_id="sync_from_subdirectory_and_allow_overwrite_and_non_recursive",
source_bucket=BUCKET_1_SRC,
source_object="subdir/",