fix: EntityType RPC update returns the updated EntityType - not an LRO.

PiperOrigin-RevId: 527003565
googleapis · Apr 25, 2023 · 8f9c714 · 8f9c714
1 parent 0b3dac9
commit 8f9c714
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 8 deletions.
diff --git a/google/cloud/aiplatform/featurestore/_entity_type.py b/google/cloud/aiplatform/featurestore/_entity_type.py
@@ -249,18 +249,15 @@ def update(
             self,
         )
 
-        update_entity_type_lro = self.api_client.update_entity_type(
+        updated_entity_type = self.api_client.update_entity_type(
             entity_type=gapic_entity_type,
             update_mask=update_mask,
             metadata=request_metadata,
             timeout=update_request_timeout,
         )
 
-        _LOGGER.log_action_started_against_resource_with_lro(
-            "Update", "entityType", self.__class__, update_entity_type_lro
-        )
-
-        update_entity_type_lro.result()
+        # Update underlying resource with response data.
+        self._gca_resource = updated_entity_type
 
         _LOGGER.log_action_completed_against_resource("entityType", "updated", self)
 

diff --git a/tests/system/aiplatform/test_featurestore.py b/tests/system/aiplatform/test_featurestore.py
@@ -36,6 +36,7 @@
 _TEST_FEATURESTORE_ID = "movie_prediction"
 _TEST_USER_ENTITY_TYPE_ID = "users"
 _TEST_MOVIE_ENTITY_TYPE_ID = "movies"
+_TEST_MOVIE_ENTITY_TYPE_UPDATE_LABELS = {"my_key_update": "my_value_update"}
 
 _TEST_USER_AGE_FEATURE_ID = "age"
 _TEST_USER_GENDER_FEATURE_ID = "gender"
@@ -129,6 +130,15 @@ def test_create_get_list_entity_types(self, shared_state):
             entity_type.resource_name for entity_type in list_entity_types
         ]
 
+        # Update information about the movie entity type.
+        assert movie_entity_type.labels != _TEST_MOVIE_ENTITY_TYPE_UPDATE_LABELS
+
+        movie_entity_type.update(
+            labels=_TEST_MOVIE_ENTITY_TYPE_UPDATE_LABELS,
+        )
+
+        assert movie_entity_type.labels == _TEST_MOVIE_ENTITY_TYPE_UPDATE_LABELS
+
     def test_create_get_list_features(self, shared_state):
 
         assert shared_state["user_entity_type"]

diff --git a/tests/unit/aiplatform/test_featurestores.py b/tests/unit/aiplatform/test_featurestores.py
@@ -591,8 +591,10 @@ def update_entity_type_mock():
     with patch.object(
         featurestore_service_client.FeaturestoreServiceClient, "update_entity_type"
     ) as update_entity_type_mock:
-        update_entity_type_lro_mock = mock.Mock(operation.Operation)
-        update_entity_type_mock.return_value = update_entity_type_lro_mock
+        update_entity_type_mock.return_value = gca_entity_type.EntityType(
+            name=_TEST_ENTITY_TYPE_NAME,
+            labels=_TEST_LABELS_UPDATE,
+        )
         yield update_entity_type_mock
 
 
@@ -2104,6 +2106,8 @@ def test_update_entity_type(self, update_entity_type_mock):
             timeout=None,
         )
 
+        assert my_entity_type.labels == _TEST_LABELS_UPDATE
+
     @pytest.mark.parametrize(
         "featurestore_name", [_TEST_FEATURESTORE_NAME, _TEST_FEATURESTORE_ID]
     )