You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by el...@apache.org on 2024/07/04 17:17:02 UTC
(airflow) branch main updated: Add tests for prevent changing main GCSToGCSOperator behavior (#40610)
This is an automated email from the ASF dual-hosted git repository.
eladkal pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 4edcac2445 Add tests for prevent changing main GCSToGCSOperator behavior (#40610)
4edcac2445 is described below
commit 4edcac24454748efcad50546df6c3587ac5148ef
Author: Maksim <ma...@google.com>
AuthorDate: Thu Jul 4 10:16:55 2024 -0700
Add tests for prevent changing main GCSToGCSOperator behavior (#40610)
---
.../google/cloud/transfers/test_gcs_to_gcs.py | 18 ++++++++++++
.../google/cloud/gcs/example_gcs_to_gcs.py | 32 ++++++++++++++++++----
2 files changed, 45 insertions(+), 5 deletions(-)
diff --git a/tests/providers/google/cloud/transfers/test_gcs_to_gcs.py b/tests/providers/google/cloud/transfers/test_gcs_to_gcs.py
index 97e16e7366..59c0b09eb0 100644
--- a/tests/providers/google/cloud/transfers/test_gcs_to_gcs.py
+++ b/tests/providers/google/cloud/transfers/test_gcs_to_gcs.py
@@ -44,6 +44,7 @@ SOURCE_OBJECT_WILDCARD_PREFIX = "*test_object"
SOURCE_OBJECT_WILDCARD_SUFFIX = "test_object*"
SOURCE_OBJECT_WILDCARD_MIDDLE = "test*object"
SOURCE_OBJECT_WILDCARD_FILENAME = "test_object*.txt"
+SOURCE_OBJECT_WILDCARD_NESTED = "test_object/file1.txt"
SOURCE_OBJECT_NO_WILDCARD = "test_object.txt"
SOURCE_OBJECT_MULTIPLE_WILDCARDS = "csv/*/test_*.csv"
DESTINATION_BUCKET = "archive"
@@ -272,6 +273,23 @@ class TestGoogleCloudStorageToCloudStorageOperator:
]
mock_hook.return_value.rewrite.assert_has_calls(mock_calls_empty)
+ @mock.patch("airflow.providers.google.cloud.transfers.gcs_to_gcs.GCSHook")
+ def test_execute_wildcard_with_destination_object_rename(self, mock_hook):
+ mock_hook.return_value.list.return_value = SOURCE_OBJECTS_SINGLE_FILE
+ operator = GCSToGCSOperator(
+ task_id=TASK_ID,
+ source_bucket=TEST_BUCKET,
+ source_object=SOURCE_OBJECT_WILDCARD_NESTED,
+ destination_bucket=DESTINATION_BUCKET,
+ destination_object="test_file.txt",
+ )
+
+ operator.execute(None)
+ mock_calls = [
+ mock.call(TEST_BUCKET, "test_object/file1.txt", DESTINATION_BUCKET, "test_file.txt"),
+ ]
+ mock_hook.return_value.rewrite.assert_has_calls(mock_calls)
+
@mock.patch("airflow.providers.google.cloud.transfers.gcs_to_gcs.GCSHook")
def test_execute_last_modified_time(self, mock_hook):
mock_hook.return_value.list.return_value = SOURCE_OBJECTS_LIST[:-1]
diff --git a/tests/system/providers/google/cloud/gcs/example_gcs_to_gcs.py b/tests/system/providers/google/cloud/gcs/example_gcs_to_gcs.py
index 746f877678..85fb95e7a8 100644
--- a/tests/system/providers/google/cloud/gcs/example_gcs_to_gcs.py
+++ b/tests/system/providers/google/cloud/gcs/example_gcs_to_gcs.py
@@ -30,9 +30,11 @@ from airflow.decorators import task
from airflow.models.baseoperator import chain
from airflow.models.dag import DAG
from airflow.operators.bash import BashOperator
+from airflow.operators.python import PythonOperator
from airflow.providers.google.cloud.operators.gcs import (
GCSCreateBucketOperator,
GCSDeleteBucketOperator,
+ GCSListObjectsOperator,
GCSSynchronizeBucketsOperator,
)
from airflow.providers.google.cloud.transfers.gcs_to_gcs import GCSToGCSOperator
@@ -52,6 +54,15 @@ HOME = "/home/airflow/gcs"
PREFIX = f"{HOME}/data/{DAG_ID}_{ENV_ID}/"
+def _assert_copied_files_exist(ti):
+ objects = ti.xcom_pull(task_ids=["list_objects"], key="return_value")[0]
+
+ assert PREFIX + OBJECT_1 in objects
+ assert f"{PREFIX}subdir/{OBJECT_1}" in objects
+ assert f"backup/{OBJECT_1}" in objects
+ assert f"backup_{OBJECT_1}" in objects
+
+
with DAG(
DAG_ID,
schedule="@once",
@@ -159,7 +170,7 @@ with DAG(
copy_single_file = GCSToGCSOperator(
task_id="copy_single_gcs_file",
source_bucket=BUCKET_NAME_SRC,
- source_object=OBJECT_1,
+ source_object=PREFIX + OBJECT_1,
destination_bucket=BUCKET_NAME_DST, # If not supplied the source_bucket value will be used
destination_object="backup_" + OBJECT_1, # If not supplied the source_object value will be used
exact_match=True,
@@ -170,7 +181,7 @@ with DAG(
copy_files = GCSToGCSOperator(
task_id="copy_files",
source_bucket=BUCKET_NAME_SRC,
- source_object="subdir/",
+ source_object=PREFIX + "subdir/",
destination_bucket=BUCKET_NAME_DST,
destination_object="backup/",
)
@@ -191,7 +202,10 @@ with DAG(
copy_files_with_list = GCSToGCSOperator(
task_id="copy_files_with_list",
source_bucket=BUCKET_NAME_SRC,
- source_objects=[OBJECT_1, OBJECT_2], # Instead of files each element could be a wildcard expression
+ source_objects=[
+ PREFIX + OBJECT_1,
+ PREFIX + OBJECT_2,
+ ], # Instead of files each element could be a wildcard expression
destination_bucket=BUCKET_NAME_DST,
destination_object="backup/",
)
@@ -201,7 +215,7 @@ with DAG(
move_single_file = GCSToGCSOperator(
task_id="move_single_file",
source_bucket=BUCKET_NAME_SRC,
- source_object=OBJECT_1,
+ source_object=PREFIX + OBJECT_1,
destination_bucket=BUCKET_NAME_DST,
destination_object="backup_" + OBJECT_1,
exact_match=True,
@@ -213,12 +227,18 @@ with DAG(
move_files_with_list = GCSToGCSOperator(
task_id="move_files_with_list",
source_bucket=BUCKET_NAME_SRC,
- source_objects=[OBJECT_1, OBJECT_2],
+ source_objects=[PREFIX + OBJECT_1, PREFIX + OBJECT_2],
destination_bucket=BUCKET_NAME_DST,
destination_object="backup/",
)
# [END howto_operator_gcs_to_gcs_list_move]
+ list_objects = GCSListObjectsOperator(task_id="list_objects", bucket=BUCKET_NAME_DST)
+ assert_copied_files_exist = PythonOperator(
+ task_id="assert_copied_files_exist",
+ python_callable=_assert_copied_files_exist,
+ )
+
delete_bucket_src = GCSDeleteBucketOperator(
task_id="delete_bucket_src", bucket_name=BUCKET_NAME_SRC, trigger_rule=TriggerRule.ALL_DONE
)
@@ -248,6 +268,8 @@ with DAG(
copy_files_with_list,
move_single_file,
move_files_with_list,
+ list_objects,
+ assert_copied_files_exist,
# TEST TEARDOWN
[delete_bucket_src, delete_bucket_dst, delete_work_dir(create_workdir_task)],
)