fix: change endpoint update method to return resource (#1409)

* fix: change endpoint update method to return resource * fix: update unit tests for endpoint.update

Author: sararob

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

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):