GCSDeleteObjectsOperator empty prefix bug fix (#24353)

apache · Jun 19, 2022 · e7a1c50 · e7a1c50
1 parent dd35fda
commit e7a1c50
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 5 deletions.
diff --git a/airflow/providers/google/cloud/operators/gcs.py b/airflow/providers/google/cloud/operators/gcs.py
@@ -305,8 +305,11 @@ def __init__(
  self.delegate_to = delegate_to
  self.impersonation_chain = impersonation_chain
 
- if not objects and not prefix:
- raise ValueError("Either object or prefix should be set. Both are None")
+ if objects is None and prefix is None:
+ err_message = "(Task {task_id}) Either object or prefix should be set. Both are None.".format(
+ **kwargs
+ )
+ raise ValueError(err_message)
 
  super().__init__(**kwargs)
 

diff --git a/tests/providers/google/cloud/operators/test_gcs.py b/tests/providers/google/cloud/operators/test_gcs.py
@@ -38,7 +38,7 @@
 TEST_PROJECT = "test-project"
 DELIMITER = ".csv"
 PREFIX = "TEST"
-MOCK_FILES = ["TEST1.csv", "TEST2.csv", "TEST3.csv"]
+MOCK_FILES = ["TEST1.csv", "TEST2.csv", "TEST3.csv", "OTHERTEST1.csv"]
 TEST_OBJECT = "dir1/test-object"
 LOCAL_FILE_PATH = "/home/airflow/gcp/test-object"
 IMPERSONATION_CHAIN = ["ACCOUNT_1", "ACCOUNT_2", "ACCOUNT_3"]
@@ -108,7 +108,7 @@ def test_object_create_acl(self, mock_hook):
  )
 
 
-class TestGoogleCloudStorageDeleteOperator(unittest.TestCase):
+class TestGCSDeleteObjectsOperator(unittest.TestCase):
  @mock.patch("airflow.providers.google.cloud.operators.gcs.GCSHook")
  def test_delete_objects(self, mock_hook):
  operator = GCSDeleteObjectsOperator(task_id=TASK_ID, bucket_name=TEST_BUCKET, objects=MOCK_FILES[0:2])
@@ -125,7 +125,7 @@ def test_delete_objects(self, mock_hook):
 
  @mock.patch("airflow.providers.google.cloud.operators.gcs.GCSHook")
  def test_delete_prefix(self, mock_hook):
- mock_hook.return_value.list.return_value = MOCK_FILES[1:3]
+ mock_hook.return_value.list.return_value = MOCK_FILES[1:4]
  operator = GCSDeleteObjectsOperator(task_id=TASK_ID, bucket_name=TEST_BUCKET, prefix=PREFIX)
 
  operator.execute(None)
@@ -138,6 +138,23 @@ def test_delete_prefix(self, mock_hook):
  any_order=True,
  )
 
+ @mock.patch("airflow.providers.google.cloud.operators.gcs.GCSHook")
+ def test_delete_prefix_as_empty_string(self, mock_hook):
+ mock_hook.return_value.list.return_value = MOCK_FILES[0:4]
+ operator = GCSDeleteObjectsOperator(task_id=TASK_ID, bucket_name=TEST_BUCKET, prefix="")
+
+ operator.execute(None)
+ mock_hook.return_value.list.assert_called_once_with(bucket_name=TEST_BUCKET, prefix="")
+ mock_hook.return_value.delete.assert_has_calls(
+ calls=[
+ mock.call(bucket_name=TEST_BUCKET, object_name=MOCK_FILES[0]),
+ mock.call(bucket_name=TEST_BUCKET, object_name=MOCK_FILES[1]),
+ mock.call(bucket_name=TEST_BUCKET, object_name=MOCK_FILES[2]),
+ mock.call(bucket_name=TEST_BUCKET, object_name=MOCK_FILES[3]),
+ ],
+ any_order=True,
+ )
+
 
 class TestGoogleCloudStorageListOperator(unittest.TestCase):
  @mock.patch("airflow.providers.google.cloud.operators.gcs.GCSHook")