Mock out all of the google auth methods by moving them to helpers

Author: pselden

litellm/llms/vertex_ai/vertex_llm_base.py

@@ -40,13 +40,7 @@ def get_vertex_region(self, vertex_region: Optional[str]) -> str:
     def load_auth(
         self, credentials: Optional[VERTEX_CREDENTIALS_TYPES], project_id: Optional[str]
     ) -> Tuple[Any, str]:
-        import google.auth as google_auth
-        from google.auth import identity_pool
-
         if credentials is not None:
-            import google.oauth2.service_account
-            import google.oauth2.credentials as google_oauth_credentials
-            
             if isinstance(credentials, str):
                 verbose_logger.debug(
                     "Vertex: Loading vertex credentials from %s", credentials
@@ -78,27 +72,25 @@ def load_auth(
 
             # Check if the JSON object contains Workload Identity Federation configuration
             if "type" in json_obj and json_obj["type"] == "external_account":
-                creds = identity_pool.Credentials.from_info(json_obj)
+                creds = self._credentials_from_identity_pool(json_obj)
             # Check if the JSON object contains Authorized User configuration (via gcloud auth application-default login)
             elif "type" in json_obj and json_obj["type"] == "authorized_user":
-                creds = google_oauth_credentials.Credentials.from_authorized_user_info(
+                creds = self._credentials_from_authorized_user(
                     json_obj,
                     scopes=["https://www.googleapis.com/auth/cloud-platform"],
                 )
                 if project_id is None:
                     project_id = creds.quota_project_id # authorized user credentials don't have a project_id, only quota_project_id
             else:
-                creds = (
-                    google.oauth2.service_account.Credentials.from_service_account_info(
-                        json_obj,
-                        scopes=["https://www.googleapis.com/auth/cloud-platform"],
-                    )
+                creds = self._credentials_from_service_account(
+                    json_obj,
+                    scopes=["https://www.googleapis.com/auth/cloud-platform"],
                 )
 
             if project_id is None:
                 project_id = getattr(creds, "project_id", None)
         else:
-            creds, creds_project_id = google_auth.default(
+            creds, creds_project_id = self._credentials_from_default_auth(
                 quota_project_id=project_id,
                 scopes=["https://www.googleapis.com/auth/cloud-platform"],
             )
@@ -116,6 +108,24 @@ def load_auth(
             )
 
         return creds, project_id
+    
+    # Google Auth Helpers -- extracted for mocking purposes in tests
+    def _credentials_from_identity_pool(self, json_obj):
+        from google.auth import identity_pool
+        return identity_pool.Credentials.from_info(json_obj)
+
+    def _credentials_from_authorized_user(self, json_obj, scopes):
+        import google.oauth2.credentials
+        return google.oauth2.credentials.Credentials.from_authorized_user_info(json_obj, scopes=scopes)
+
+    def _credentials_from_service_account(self, json_obj, scopes):
+        import google.oauth2.service_account
+        return google.oauth2.service_account.Credentials.from_service_account_info(json_obj, scopes=scopes)
+
+    def _credentials_from_default_auth(self, quota_project_id, scopes):
+        import google.auth as google_auth
+        return google_auth.default(quota_project_id=quota_project_id, scopes=scopes)
+
 
     def refresh_auth(self, credentials: Any) -> None:
         from google.auth.transport.requests import (

poetry.lock

@@ -58,7 +58,6 @@ redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.
 mcp = {version = "1.5.0", optional = true, python = ">=3.10"}
 litellm-proxy-extras = {version = "0.1.21", optional = true}
 rich = {version = "13.7.1", optional = true}
-google-auth = {version = "^2.40.1", optional = true}
 litellm-enterprise = {version = "0.1.4", optional = true}
 
 [tool.poetry.extras]

pyproject.toml

@@ -175,12 +175,13 @@ async def test_gemini_credentials(self, is_async):
         assert token == ""
         assert project == ""
 
-
     @pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"])
     @pytest.mark.asyncio
     async def test_authorized_user_credentials(self, is_async):
         vertex_base = VertexBase()
 
+        quota_project_id = "test-project"
+
         credentials = {
             "account": "",
             "client_id": "fake-client-id",
@@ -191,7 +192,15 @@ async def test_authorized_user_credentials(self, is_async):
             "universe_domain": "googleapis.com"
         }
 
-        with patch.object(vertex_base, "refresh_auth") as mock_refresh:
+        mock_creds = MagicMock()
+        mock_creds.token = "token-1"
+        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:
             def mock_refresh_impl(creds):
                 creds.token = "refreshed-token"
 
@@ -207,20 +216,21 @@ def mock_refresh_impl(creds):
                     credentials=credentials, project_id=None, custom_llm_provider="vertex_ai"
                 )
 
-
+            assert mock_credentials_from_authorized_user.called
             assert token == "refreshed-token"
-            assert project == "test-project"
+            assert project == quota_project_id
 
             # 2. Test that authorized_user-style credentials are correctly handled and uses passed in project_id
+            not_quota_project_id = "new-project"
             if is_async:
                 token, project = await vertex_base._ensure_access_token_async(
-                    credentials=credentials, project_id="new-project", 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="new-project", custom_llm_provider="vertex_ai"
+                    credentials=credentials, project_id=not_quota_project_id, custom_llm_provider="vertex_ai"
                 )
 
 
             assert token == "refreshed-token"
-            assert project == "new-project"
+            assert project == not_quota_project_id