chore: fix unit test assertions

PiperOrigin-RevId: 711773190
googleapis · Feb 1, 2025 · d17c88c · d17c88c
1 parent 58fbabd
commit d17c88c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/tests/unit/aiplatform/test_endpoints.py b/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)
 

diff --git a/tests/unit/aiplatform/test_utils.py b/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"