diff --git a/google/cloud/aiplatform/featurestore/_entity_type.py b/google/cloud/aiplatform/featurestore/_entity_type.py index 859e5fc7ce..fee48a05cb 100644 --- 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 index 7a05f3c15d..cf875bac85 100644 --- 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 index 25a939c795..3eccdb401f 100644 --- 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] )