From 28c06533150d714a0d9fa483d87fcc61989a4f38 Mon Sep 17 00:00:00 2001
From: Morgan Du <morgandu@google.com>
Date: Tue, 15 Mar 2022 18:13:20 -0700
Subject: [PATCH] fix: add self.wait() in operations after optional_sync
 supported resource creation

---
 .../matching_engine/matching_engine_index_endpoint.py      | 7 +++++++
 google/cloud/aiplatform/models.py                          | 4 ++++
 2 files changed, 11 insertions(+)

diff --git a/google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py b/google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py
index c585906ea0..c86cff77c2 100644
--- 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
index aaef433aeb..5029f4db6d 100644
--- 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.")