update sync usage in deleting child resources

googleapis · Nov 30, 2021 · 8742825 · 8742825
1 parent e50e763
commit 8742825
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 10 deletions.
diff --git a/google/cloud/aiplatform/_featurestores/entity_type.py b/google/cloud/aiplatform/_featurestores/entity_type.py
@@ -361,9 +361,7 @@ def list_features(
         )
 
     @base.optional_sync()
-    def delete_features(
-        self, feature_ids: List[str], sync: Optional[bool] = True,
-    ) -> None:
+    def delete_features(self, feature_ids: List[str], sync: bool = True,) -> None:
         """Deletes feature resources in this EntityType given their feature IDs.
         WARNING: This deletion is permanent.
 
@@ -375,9 +373,11 @@ def delete_features(
                 will be executed in concurrent Future and any downstream object will
                 be immediately returned and synced when the Future has completed.
         """
+        features = []
         for feature_id in feature_ids:
             feature = self.get_feature(feature_id=feature_id)
-            feature.delete(sync=sync)
+            feature.delete(sync=False)
+            features.append(feature)
 
-        if not sync:
+        for feature in features:
             feature.wait()
diff --git a/google/cloud/aiplatform/_featurestores/featurestore.py b/google/cloud/aiplatform/_featurestores/featurestore.py
@@ -301,9 +301,9 @@ def list_entity_types(
 
     @base.optional_sync()
     def delete_entity_types(
-        self, entity_type_ids: List[str], sync: Optional[bool] = True,
+        self, entity_type_ids: List[str], sync: bool = True,
     ) -> None:
-        """Deletes entity_type resources in this Featurestre given their entity_type IDs.
+        """Deletes entity_type resources in this Featurestore given their entity_type IDs.
         WARNING: This deletion is permanent.
 
         Args:
@@ -314,9 +314,11 @@ def delete_entity_types(
                 will be executed in concurrent Future and any downstream object will
                 be immediately returned and synced when the Future has completed.
         """
+        entity_types = []
         for entity_type_id in entity_type_ids:
             entity_type = self.get_entity_type(entity_type_id=entity_type_id)
-            entity_type.delete(sync=sync)
+            entity_type.delete(sync=False)
+            entity_types.append(entity_type)
 
-        if not sync:
+        for entity_type in entity_types:
             entity_type.wait()
diff --git a/tests/system/aiplatform/test_featurestore.py b/tests/system/aiplatform/test_featurestore.py
@@ -22,7 +22,7 @@
 
 class TestFeaturestore(e2e_base.TestEndToEnd):
 
-    _temp_prefix = "temp-vertex-sdk-e2e-test"
+    _temp_prefix = "temp-vertex-sdk-e2e-feature-store-test"
 
     def test_create_and_get_featurestore(self, shared_state):