diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 445e34966a9..fd16427fabe 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -37,6 +37,15 @@ else: GoogleCredentialsObject = Any +def _import_token_state() -> "type[TokenState]": + try: + from google.auth.credentials import TokenState + except ImportError: + raise ImportError(GOOGLE_IMPORT_ERROR_MESSAGE) + + return TokenState + + class VertexBase: def __init__(self) -> None: super().__init__() @@ -397,7 +406,7 @@ class VertexBase: Look up cached credentials and return (token, project_id) if the token is FRESH. Returns None if not cached or not fresh. """ - from google.auth.credentials import TokenState + TokenState: Final = _import_token_state() creds, cached_project_id = self._unpack_cached_credentials(credential_cache_key) if ( @@ -423,7 +432,7 @@ class VertexBase: credentials object so the caller can schedule a background refresh without holding the per-key async lock. """ - from google.auth.credentials import TokenState + TokenState: Final = _import_token_state() creds, cached_project_id = self._unpack_cached_credentials(credential_cache_key) if creds is None: @@ -457,7 +466,7 @@ class VertexBase: Falls back to expired/valid checks if token_state is unavailable (e.g. older google-auth versions or mock objects in tests). """ - from google.auth.credentials import TokenState as _TokenState + _TokenState: Final = _import_token_state() token_state: Final = getattr(credentials, "token_state", None) if isinstance(token_state, _TokenState): @@ -978,7 +987,7 @@ class VertexBase: only one coroutine refreshes while others wait on the lock. Uses native async refresh for service_account and authorized_user credentials. """ - from google.auth.credentials import TokenState + TokenState: Final = _import_token_state() cache_credentials: Final = json.dumps(credentials) if isinstance(credentials, dict) else credentials credential_cache_key: Final = (cache_credentials, project_id) diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py index 18fc239b7c6..75d9086ced1 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py @@ -2088,3 +2088,33 @@ class TestVertexBase: assert token == "cached-token" assert not mock_get_lock.called, "Fast path should not acquire lock" + + @pytest.mark.asyncio + async def test_missing_google_auth_raises_actionable_import_error(self): + """Test that a missing google-auth surfaces the install hint, not a raw ModuleNotFoundError""" + vertex_base = VertexBase() + + with patch.dict("sys.modules", {"google.auth.credentials": None}): + with pytest.raises(ImportError, match="Google Cloud SDK not found"): + await vertex_base._ensure_access_token_async( + credentials={"type": "service_account", "project_id": "project-1"}, + project_id="project-1", + custom_llm_provider="vertex_ai", + ) + + @pytest.mark.parametrize( + "call_helper", + [ + lambda vb: vb._try_get_cached_token(("key", None), "project-1"), + lambda vb: vb._try_get_usable_cached_token(("key", None), "project-1"), + lambda vb: vb._get_token_state(MagicMock()), + ], + ids=["try_get_cached_token", "try_get_usable_cached_token", "get_token_state"], + ) + def test_token_state_helpers_raise_actionable_import_error(self, call_helper): + """Test that every TokenState call site surfaces the install hint""" + vertex_base = VertexBase() + + with patch.dict("sys.modules", {"google.auth.credentials": None}): + with pytest.raises(ImportError, match="Google Cloud SDK not found"): + call_helper(vertex_base)