Merge branch 'main' into artifact_registry

googleapis · Jun 7, 2022 · d70b58f · d70b58f
2 parents 787ac03 + 7a54637
commit d70b58f
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 11 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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):