fix: add self.wait() in operations after optional_sync supported resource creation (#1083)

morgandu · parthea · web-flow · commit 79aeec138006 · 2022-03-16T10:55:04.000-07:00
Co-authored-by: Anthonios Partheniou &lt;partheniou@google.com&gt;
diff --git a/google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py b/google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py
@@ -535,6 +535,8 @@ def deploy_index(
             MatchingEngineIndexEndpoint - IndexEndpoint resource object
         """
 
+        self.wait()
+
         _LOGGER.log_action_start_against_resource(
             "Deploying index", "index_endpoint", self,
         )
@@ -591,6 +593,8 @@ def undeploy_index(
             MatchingEngineIndexEndpoint - IndexEndpoint resource object
         """
 
+        self.wait()
+
         _LOGGER.log_action_start_against_resource(
             "Undeploying index", "index_endpoint", self,
         )
@@ -652,6 +656,8 @@ def mutate_deployed_index(
                 Optional. Strings which should be sent along with the request as metadata.
         """
 
+        self.wait()
+
         _LOGGER.log_action_start_against_resource(
             "Mutating index", "index_endpoint", self,
         )
@@ -691,6 +697,7 @@ def deployed_indexes(
         Returns:
             List[gca_matching_engine_index_endpoint.DeployedIndex] - Deployed indexes
         """
+        self._assert_gca_resource_is_available()
         return self._gca_resource.deployed_indexes
 
     @base.optional_sync()
diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py
@@ -1575,6 +1575,8 @@ def update(
             ValueError: If `labels` is not the correct format.
         """
 
+        self.wait()
+
         current_model_proto = self.gca_resource
         copied_model_proto = current_model_proto.__class__(current_model_proto)
 
@@ -2496,6 +2498,8 @@ def export_model(
             ValueError: If invalid arguments or export formats are provided.
         """
 
+        self.wait()
+
         # Model does not support exporting
         if not self.supported_export_formats:
             raise ValueError(f"The model `{self.resource_name}` is not exportable.")
