fix tests

Author: twishabansal

packages/toolbox-core/tests/test_auth_methods.py

@@ -44,6 +44,9 @@ def reset_cache():
     auth_methods._token_cache = original_cache
 
 
+
+
+
 @pytest.mark.asyncio
 class TestAsyncAuthMethods:
     """Tests for asynchronous Google ID token fetching."""
@@ -61,7 +64,8 @@ async def test_aget_google_id_token_success_first_call(
         mock_fetch.return_value = MOCK_ID_TOKEN
         mock_verify.return_value = {"exp": MOCK_EXPIRY_TIMESTAMP}
 
-        token = await auth_methods.aget_google_id_token(MOCK_AUDIENCE)
+        token_getter = auth_methods.aget_google_id_token(MOCK_AUDIENCE)
+        token = await token_getter()
 
         mock_default.assert_called_once()
         mock_fetch.assert_called_once_with(ANY, MOCK_AUDIENCE)
@@ -72,15 +76,14 @@ async def test_aget_google_id_token_success_first_call(
     @patch("toolbox_core.auth_methods.google.auth.default")
     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
         auth_methods._token_cache["expires_at"] = auth_methods.datetime.now(
             auth_methods.timezone.utc
         ) + auth_methods.timedelta(hours=1)
 
-        token = await auth_methods.aget_google_id_token(MOCK_AUDIENCE)
+        token_getter = auth_methods.aget_google_id_token(MOCK_AUDIENCE)
+        token = await token_getter()
 
-        # 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}"
 
@@ -97,14 +100,13 @@ async def test_aget_google_id_token_refreshes_expired_cache(
         auth_methods._token_cache["token"] = "expired_token"
         auth_methods._token_cache["expires_at"] = auth_methods.datetime.now(
             auth_methods.timezone.utc
-        ) - auth_methods.timedelta(
-            seconds=100
-        )  # Expired
+        ) - auth_methods.timedelta(seconds=100)
 
         mock_fetch.return_value = MOCK_ID_TOKEN
         mock_verify.return_value = {"exp": MOCK_EXPIRY_TIMESTAMP}
 
-        token = await auth_methods.aget_google_id_token(MOCK_AUDIENCE)
+        token_getter = auth_methods.aget_google_id_token(MOCK_AUDIENCE)
+        token = await token_getter()
 
         mock_default.assert_called_once()
         mock_fetch.assert_called_once()
@@ -122,8 +124,8 @@ async def test_aget_raises_if_no_audience_and_no_local_token(
         """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()
+            token_getter = auth_methods.aget_google_id_token()
+            await token_getter()
 
         mock_default.assert_called_once()
         mock_fetch.assert_not_called()
@@ -149,7 +151,8 @@ def test_get_google_id_token_success_local_creds(
         mock_request_instance = MagicMock()
         mock_request.return_value = mock_request_instance
 
-        token = auth_methods.get_google_id_token(MOCK_AUDIENCE)
+        token_getter = auth_methods.get_google_id_token(MOCK_AUDIENCE)
+        token = token_getter()
 
         mock_default.assert_called_once_with()
         mock_session.assert_called_once_with(mock_creds)
@@ -167,7 +170,8 @@ def test_get_google_id_token_success_uses_cache(self, mock_default):
             auth_methods.timezone.utc
         ) + auth_methods.timedelta(hours=1)
 
-        token = auth_methods.get_google_id_token(MOCK_AUDIENCE)
+        token_getter = auth_methods.get_google_id_token(MOCK_AUDIENCE)
+        token = token_getter()
 
         mock_default.assert_not_called()
         assert token == f"{auth_methods.BEARER_TOKEN_PREFIX}{MOCK_ID_TOKEN}"
@@ -185,7 +189,7 @@ def test_get_google_id_token_fetch_failure(
         mock_fetch.side_effect = GoogleAuthError("Fetch failed")
 
         with pytest.raises(GoogleAuthError, match="Fetch failed"):
-            auth_methods.get_google_id_token(MOCK_AUDIENCE)
+            auth_methods.get_google_id_token(MOCK_AUDIENCE)()
 
         assert auth_methods._token_cache["token"] is None
         mock_default.assert_called_once()
@@ -208,15 +212,13 @@ def test_get_google_id_token_validation_failure(
         with pytest.raises(
             ValueError, match="Failed to validate and cache the new token"
         ):
-            auth_methods.get_google_id_token(MOCK_AUDIENCE)
+            auth_methods.get_google_id_token(MOCK_AUDIENCE)()
 
-        # 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(
@@ -225,8 +227,7 @@ def test_get_raises_if_no_audience_and_no_local_token(
         """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()
+            auth_methods.get_google_id_token()()
 
         mock_default.assert_called_once()
-        mock_fetch.assert_not_called()
+        mock_fetch.assert_not_called()