add optional arg

Author: twishabansal

packages/toolbox-core/src/toolbox_core/auth_methods.py

@@ -36,7 +36,7 @@
 
 import asyncio
 from datetime import datetime, timedelta, timezone
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 import google.auth
 from google.auth.exceptions import GoogleAuthError
@@ -94,7 +94,7 @@ def _update_cache(new_token: str) -> None:
         raise ValueError(f"Failed to validate and cache the new token: {e}") from e
 
 
-def get_google_id_token(audience: str) -> str:
+def get_google_id_token(audience: Optional[str] = None) -> str:
     """
     Synchronously fetches a Google ID token for a specific audience.
     This function uses Application Default Credentials for local systems
@@ -125,8 +125,12 @@ def get_google_id_token(audience: str) -> str:
         if new_id_token:
             _update_cache(new_id_token)
             return BEARER_TOKEN_PREFIX + new_id_token
+        
+    if audience is None:
+        raise Exception('You are not authenticating using User Credentials.'
+                        ' Please set the audience string to the Toolbox service URL to get the Google ID token.')
 
-    # Get credentials for Google Cloud environments
+    # Get credentials for Google Cloud environments or for service account key files
     try:
         request = Request()
         new_token = id_token.fetch_id_token(request, audience)
@@ -139,7 +143,7 @@ def get_google_id_token(audience: str) -> str:
         ) from e
 
 
-async def aget_google_id_token(audience: str) -> str:
+async def aget_google_id_token(audience: Optional[str] = None) -> str:
     """
     Asynchronously fetches a Google ID token for a specific audience.
     This function uses Application Default Credentials for local systems

packages/toolbox-core/tests/test_auth_methods.py

@@ -69,9 +69,8 @@ async def test_aget_google_id_token_success_first_call(
         assert auth_methods._token_cache["token"] == MOCK_ID_TOKEN
         assert auth_methods._token_cache["expires_at"] == MOCK_EXPIRY_DATETIME
 
-    @pytest.mark.asyncio
     @patch("toolbox_core.auth_methods.google.auth.default")
-    async def test_aget_google_id_token_success_uses_cache(self, mock_get_token):
+    async def test_aget_google_id_token_success_uses_cache(self, mock_default):
         """Tests that subsequent calls use the cached token if valid."""
         # Prime the cache with a valid token
         auth_methods._token_cache["token"] = MOCK_ID_TOKEN
@@ -81,8 +80,8 @@ async def test_aget_google_id_token_success_uses_cache(self, mock_get_token):
 
         token = await auth_methods.aget_google_id_token(MOCK_AUDIENCE)
 
-        # The underlying sync function should not be called if cache is valid
-        mock_get_token.assert_not_called()
+        # The underlying auth function should not be called if cache is valid
+        mock_default.assert_not_called()
         assert token == f"{auth_methods.BEARER_TOKEN_PREFIX}{MOCK_ID_TOKEN}"
 
     @patch("toolbox_core.auth_methods.id_token.verify_oauth2_token")
@@ -112,6 +111,23 @@ async def test_aget_google_id_token_refreshes_expired_cache(
         assert token == f"{auth_methods.BEARER_TOKEN_PREFIX}{MOCK_ID_TOKEN}"
         assert auth_methods._token_cache["token"] == MOCK_ID_TOKEN
 
+    @patch("toolbox_core.auth_methods.id_token.fetch_id_token")
+    @patch(
+        "toolbox_core.auth_methods.google.auth.default",
+        return_value=(MagicMock(id_token=None), MOCK_PROJECT_ID),
+    )
+    async def test_aget_raises_if_no_audience_and_no_local_token(
+        self, mock_default, mock_fetch
+    ):
+        """Tests that the async function propagates the missing audience exception."""
+        error_msg = "You are not authenticating using User Credentials."
+        with pytest.raises(Exception, match=error_msg):
+            # Call without audience to trigger the error path
+            await auth_methods.aget_google_id_token()
+
+        mock_default.assert_called_once()
+        mock_fetch.assert_not_called()
+
 
 class TestSyncAuthMethods:
     """Tests for synchronous Google ID token fetching."""
@@ -196,3 +212,21 @@ def test_get_google_id_token_validation_failure(
 
         # Verify cache is cleared on validation failure
         assert auth_methods._token_cache["token"] is None
+
+    @patch("toolbox_core.auth_methods.id_token.fetch_id_token")
+    @patch(
+        "toolbox_core.auth_methods.google.auth.default",
+        # Simulate credentials that DON'T have a pre-existing id_token
+        return_value=(MagicMock(id_token=None), MOCK_PROJECT_ID),
+    )
+    def test_get_raises_if_no_audience_and_no_local_token(
+        self, mock_default, mock_fetch
+    ):
+        """Tests exception is raised if audience is required but not provided."""
+        error_msg = "You are not authenticating using User Credentials."
+        with pytest.raises(Exception, match=error_msg):
+            # Call without an audience to trigger the error path
+            auth_methods.get_google_id_token()
+
+        mock_default.assert_called_once()
+        mock_fetch.assert_not_called()