chore: fix unit test assertions

PiperOrigin-RevId: 723678874

Author: yinghsienwu

setup.py

@@ -197,6 +197,7 @@
     full_extra_require
     + profiler_extra_require
     + tokenization_testing_extra_require
+    + vizier_extra_require
     + [
         # aiohttp is required for async rest tests (need google-auth[aiohttp],
         # but can't specify extras in constraints files)

testing/constraints-3.10.txt

@@ -14,4 +14,5 @@ pytest-xdist==3.3.1 # Pinned to unbreak unit tests
 ray==2.4.0 # Pinned until 2.9.3 is verified for Ray tests
 ipython==8.22.2 # Pinned to unbreak TypeAliasType import error
 scikit-learn!=1.4.1.post1 # Pin to unbreak test_sklearn (b/332610038)
-requests==2.31.0 # Pinned to unbreak http+docker error (b/342669351)
+requests==2.31.0 # Pinned to unbreak http+docker error (b/342669351)
+google-vizier==0.1.21

testing/constraints-3.11.txt

@@ -11,3 +11,4 @@ packaging==24.1 # Increased to unbreak canonicalize_version error (b/377774673)
 pytest-xdist==3.3.1 # Pinned to unbreak unit tests
 ray==2.5.0 # Pinned until 2.9.3 is verified for Ray tests
 ipython==8.22.2 # Pinned to unbreak TypeAliasType import error
+google-vizier==0.1.21

testing/constraints-3.8.txt

@@ -12,3 +12,4 @@ packaging==24.1 # Increased to unbreak canonicalize_version error (b/377774673)
 grpcio-testing==1.34.0
 pytest-xdist==3.3.1 # Pinned to unbreak unit tests
 ray==2.4.0 # Pinned until 2.9.3 is verified for Ray tests
+google-vizier==0.1.21

testing/constraints-3.9.txt

@@ -11,3 +11,4 @@ packaging==24.1 # Increased to unbreak canonicalize_version error (b/377774673)
 grpcio-testing==1.34.0
 pytest-xdist==3.3.1 # Pinned to unbreak unit tests
 ray==2.4.0 # Pinned until 2.9.3 is verified for Ray tests
+google-vizier==0.1.21

tests/unit/aiplatform/test_endpoints.py

@@ -3922,7 +3922,7 @@ def test_psa_delete_with_force(self, sdk_undeploy_mock, delete_endpoint_mock, sy
             test_endpoint.wait()
 
         # undeploy() should not be called unless force is set to True
-        sdk_undeploy_mock.called_once_with(deployed_model_id=_TEST_ID, sync=sync)
+        sdk_undeploy_mock.assert_not_called()
 
         delete_endpoint_mock.assert_called_once_with(name=_TEST_ENDPOINT_NAME)
 

tests/unit/aiplatform/test_utils.py

@@ -559,15 +559,19 @@ def test_validate_transformations_with_column_transformations(self):
 class TestGcsUtils:
     def test_upload_to_gcs(self, json_file, mock_storage_blob_upload_from_filename):
         gcs_utils.upload_to_gcs(json_file, f"gs://{GCS_BUCKET}/{GCS_PREFIX}")
-        assert mock_storage_blob_upload_from_filename.called_once_with(json_file)
+        mock_storage_blob_upload_from_filename.assert_called_once_with(
+            filename=json_file
+        )
 
     def test_stage_local_data_in_gcs(
         self, json_file, mock_datetime, mock_storage_blob_upload_from_filename
     ):
         timestamp = EXPECTED_TIME.isoformat(sep="-", timespec="milliseconds")
         staging_gcs_dir = f"gs://{GCS_BUCKET}/{GCS_PREFIX}"
         data_uri = gcs_utils.stage_local_data_in_gcs(json_file, staging_gcs_dir)
-        assert mock_storage_blob_upload_from_filename.called_once_with(json_file)
+        mock_storage_blob_upload_from_filename.assert_called_once_with(
+            filename=json_file
+        )
         assert (
             data_uri
             == f"{staging_gcs_dir}/vertex_ai_auto_staging/{timestamp}/test.json"