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