diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py index 7d104b3112..3b4b03af3e 100644 --- a/google/cloud/aiplatform/models.py +++ b/google/cloud/aiplatform/models.py @@ -1377,19 +1377,13 @@ def update( self, ) - update_endpoint_lro = self.api_client.update_endpoint( + self._gca_resource = self.api_client.update_endpoint( endpoint=copied_endpoint_proto, update_mask=update_mask, metadata=request_metadata, timeout=update_request_timeout, ) - _LOGGER.log_action_started_against_resource_with_lro( - "Update", "endpoint", self.__class__, update_endpoint_lro - ) - - update_endpoint_lro.result() - _LOGGER.log_action_completed_against_resource("endpoint", "updated", self) return self diff --git a/setup.py b/setup.py index 7db8ad5f27..d1d01a6bde 100644 --- a/setup.py +++ b/setup.py @@ -103,8 +103,9 @@ # Until this issue is closed # https://github.com/googleapis/google-cloud-python/issues/10566 "google-api-core[grpc] >= 1.31.5, <3.0.0dev,!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0", - "proto-plus >= 1.15.0", - "packaging >= 14.3", + "proto-plus >= 1.15.0, <2.0.0dev", + "protobuf >= 3.19.0, <4.0.0dev", + "packaging >= 14.3, <22.0.0dev", "google-cloud-storage >= 1.32.0, < 3.0.0dev", "google-cloud-bigquery >= 1.15.0, < 3.0.0dev", "google-cloud-resource-manager >= 1.3.3, < 3.0.0dev", diff --git a/testing/constraints-3.6.txt b/testing/constraints-3.6.txt index 0380c54f59..acf8ecba0c 100644 --- a/testing/constraints-3.6.txt +++ b/testing/constraints-3.6.txt @@ -8,6 +8,7 @@ google-api-core==1.31.5 libcst==0.2.5 proto-plus==1.15.0 +protobuf==3.19.0 mock==4.0.2 google-cloud-storage==1.32.0 google-auth==1.25.0 # TODO: Remove when google-api-core >= 1.26.0 is required diff --git a/tests/unit/aiplatform/test_endpoints.py b/tests/unit/aiplatform/test_endpoints.py index 3e79328b34..50a6c93303 100644 --- a/tests/unit/aiplatform/test_endpoints.py +++ b/tests/unit/aiplatform/test_endpoints.py @@ -279,8 +279,11 @@ def update_endpoint_mock(): with mock.patch.object( endpoint_service_client.EndpointServiceClient, "update_endpoint" ) as update_endpoint_mock: - update_endpoint_lro_mock = mock.Mock(ga_operation.Operation) - update_endpoint_mock.return_value = update_endpoint_lro_mock + update_endpoint_mock.return_value = gca_endpoint.Endpoint( + display_name=_TEST_DISPLAY_NAME, + name=_TEST_ENDPOINT_NAME, + encryption_spec=_TEST_ENCRYPTION_SPEC, + ) yield update_endpoint_mock @@ -768,9 +771,18 @@ def test_update_endpoint(self, update_endpoint_mock): timeout=_TEST_TIMEOUT, ) + update_endpoint_mock.return_value = gca_endpoint.Endpoint( + name=_TEST_ENDPOINT_NAME, + display_name=_TEST_DISPLAY_NAME, + description=_TEST_DESCRIPTION, + labels=_TEST_LABELS, + encryption_spec=_TEST_ENCRYPTION_SPEC, + ) + @pytest.mark.usefixtures("get_endpoint_with_models_mock") def test_update_traffic_split(self, update_endpoint_mock): endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) + endpoint.update(traffic_split={_TEST_ID: 10, _TEST_ID_2: 80, _TEST_ID_3: 10}) expected_endpoint = gca_endpoint.Endpoint( @@ -788,6 +800,12 @@ def test_update_traffic_split(self, update_endpoint_mock): timeout=_TEST_TIMEOUT, ) + update_endpoint_mock.return_value = gca_endpoint.Endpoint( + display_name=_TEST_DISPLAY_NAME, + name=_TEST_ENDPOINT_NAME, + traffic_split={_TEST_ID: 10, _TEST_ID_2: 80, _TEST_ID_3: 10}, + ) + @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") @pytest.mark.parametrize("sync", [True, False]) def test_deploy(self, deploy_model_mock, sync):