From edfd8a138dc072038746de59231b908a98ee10a4 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 19 May 2025 20:26:48 -0700
Subject: [PATCH] fix: update test to reflect change

---
 .../llms/vertex_ai/test_vertex_llm_base.py    | 56 +++++++++++--------
 1 file changed, 32 insertions(+), 24 deletions(-)

diff --git a/tests/litellm/llms/vertex_ai/test_vertex_llm_base.py b/tests/litellm/llms/vertex_ai/test_vertex_llm_base.py
index e1bdb31c7dd..9704bd4294b 100644
--- a/tests/litellm/llms/vertex_ai/test_vertex_llm_base.py
+++ b/tests/litellm/llms/vertex_ai/test_vertex_llm_base.py
@@ -53,24 +53,23 @@ class TestVertexBase:
             assert project == "project-1"
             assert token == "fake-token-1"
 
-        # Test case 2: Prevent using credentials from different project
+        # Test case 2: Allow using credentials from different project
         with patch.object(
             vertex_base, "load_auth", return_value=(mock_creds, "project-1")
         ):
-            with pytest.raises(ValueError, match="Could not resolve project_id"):
-                if is_async:
-                    result = await vertex_base._ensure_access_token_async(
-                        credentials={"type": "service_account"},
-                        project_id="different-project",
-                        custom_llm_provider="vertex_ai",
-                    )
-                else:
-                    result = vertex_base._ensure_access_token(
-                        credentials={"type": "service_account"},
-                        project_id="different-project",
-                        custom_llm_provider="vertex_ai",
-                    )
-                print(f"result: {result}")
+            if is_async:
+                result = await vertex_base._ensure_access_token_async(
+                    credentials={"type": "service_account"},
+                    project_id="different-project",
+                    custom_llm_provider="vertex_ai",
+                )
+            else:
+                result = vertex_base._ensure_access_token(
+                    credentials={"type": "service_account"},
+                    project_id="different-project",
+                    custom_llm_provider="vertex_ai",
+                )
+            print(f"result: {result}")
 
     @pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"])
     @pytest.mark.asyncio
@@ -189,7 +188,7 @@ class TestVertexBase:
             "quota_project_id": "test-project",
             "refresh_token": "fake-refresh-token",
             "type": "authorized_user",
-            "universe_domain": "googleapis.com"
+            "universe_domain": "googleapis.com",
         }
 
         mock_creds = MagicMock()
@@ -197,10 +196,12 @@ class TestVertexBase:
         mock_creds.expired = False
         mock_creds.quota_project_id = quota_project_id
 
-
         with patch.object(
             vertex_base, "_credentials_from_authorized_user", return_value=mock_creds
-        ) as mock_credentials_from_authorized_user, patch.object(vertex_base, "refresh_auth") as mock_refresh:
+        ) as mock_credentials_from_authorized_user, patch.object(
+            vertex_base, "refresh_auth"
+        ) as mock_refresh:
+
             def mock_refresh_impl(creds):
                 creds.token = "refreshed-token"
 
@@ -209,11 +210,15 @@ class TestVertexBase:
             # 1. Test that authorized_user-style credentials are correctly handled and uses quota_project_id
             if is_async:
                 token, project = await vertex_base._ensure_access_token_async(
-                    credentials=credentials, project_id=None, custom_llm_provider="vertex_ai"
+                    credentials=credentials,
+                    project_id=None,
+                    custom_llm_provider="vertex_ai",
                 )
             else:
                 token, project = vertex_base._ensure_access_token(
-                    credentials=credentials, project_id=None, custom_llm_provider="vertex_ai"
+                    credentials=credentials,
+                    project_id=None,
+                    custom_llm_provider="vertex_ai",
                 )
 
             assert mock_credentials_from_authorized_user.called
@@ -224,13 +229,16 @@ class TestVertexBase:
             not_quota_project_id = "new-project"
             if is_async:
                 token, project = await vertex_base._ensure_access_token_async(
-                    credentials=credentials, project_id=not_quota_project_id, custom_llm_provider="vertex_ai"
+                    credentials=credentials,
+                    project_id=not_quota_project_id,
+                    custom_llm_provider="vertex_ai",
                 )
             else:
                 token, project = vertex_base._ensure_access_token(
-                    credentials=credentials, project_id=not_quota_project_id, custom_llm_provider="vertex_ai"
+                    credentials=credentials,
+                    project_id=not_quota_project_id,
+                    custom_llm_provider="vertex_ai",
                 )
 
-
             assert token == "refreshed-token"
-            assert project == not_quota_project_id
\ No newline at end of file
+            assert project == not_quota_project_id