Skip to content

Commit

Permalink
fix: fix feature update since no LRO is created
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 562005959
  • Loading branch information
matthew29tang authored and copybara-github committed Sep 1, 2023
1 parent 2bd494b commit 468e6e7
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 13 deletions.
11 changes: 1 addition & 10 deletions google/cloud/aiplatform/featurestore/feature.py
Original file line number Diff line number Diff line change
Expand Up @@ -238,21 +238,12 @@ def update(
self,
)

update_feature_lro = self.api_client.update_feature(
self._gca_resource = self.api_client.update_feature(
feature=gapic_feature,
update_mask=update_mask,
metadata=request_metadata,
timeout=update_request_timeout,
)

_LOGGER.log_action_started_against_resource_with_lro(
"Update", "feature", self.__class__, update_feature_lro
)

update_feature_lro.result()

_LOGGER.log_action_completed_against_resource("feature", "updated", self)

return self

@classmethod
Expand Down
11 changes: 8 additions & 3 deletions tests/unit/aiplatform/test_featurestores.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
from google.cloud.aiplatform.utils import resource_manager_utils

from google.cloud.aiplatform.utils import featurestore_utils
from google.cloud.aiplatform.featurestore.feature import Feature
from google.cloud.aiplatform.compat.services import (
featurestore_service_client,
)
Expand Down Expand Up @@ -721,8 +722,10 @@ def update_feature_mock():
with patch.object(
featurestore_service_client.FeaturestoreServiceClient, "update_feature"
) as update_feature_mock:
update_feature_lro_mock = mock.Mock(operation.Operation)
update_feature_mock.return_value = update_feature_lro_mock
update_feature_mock.return_value = gca_feature.Feature(
name=_TEST_FEATURE_NAME,
value_type=_TEST_FEATURE_VALUE_TYPE_ENUM,
)
yield update_feature_mock


Expand Down Expand Up @@ -3693,7 +3696,7 @@ def test_update_feature(self, update_feature_mock):
aiplatform.init(project=_TEST_PROJECT)

my_feature = aiplatform.Feature(feature_name=_TEST_FEATURE_NAME)
my_feature.update(
updated_feature = my_feature.update(
labels=_TEST_LABELS_UPDATE,
update_request_timeout=None,
)
Expand All @@ -3709,6 +3712,8 @@ def test_update_feature(self, update_feature_mock):
timeout=None,
)

assert isinstance(updated_feature, Feature)

@pytest.mark.parametrize(
"entity_type_name, featurestore_id",
[
Expand Down

0 comments on commit 468e6e7

Please sign in to comment.