Fix failing unit test, and change acceptable values for the env variable

Author: nbayati

google/auth/_helpers.py

@@ -19,6 +19,7 @@
 import datetime
 from email.message import Message
 import hashlib
+import os
 import json
 import logging
 import sys
@@ -287,6 +288,45 @@ def unpadded_urlsafe_b64encode(value):
     return base64.urlsafe_b64encode(value).rstrip(b"=")
 
 
+def get_bool_from_env(variable_name, default=False):
+    """Gets a boolean value from an environment variable.
+
+    The environment variable is interpreted as a boolean with the following
+    (case-insensitive) rules:
+    - "true", "1" are considered true.
+    - "false", "0" are considered false.
+
+    Args:
+        variable_name (str): The name of the environment variable.
+        default (bool): The default value if the environment variable is not
+            set.
+
+    Returns:
+        bool: The boolean value of the environment variable.
+
+    Raises:
+        google.auth.exceptions.InvalidValue: If the environment variable is
+            set to a value that can not be interpreted as a boolean.
+    """
+    value = os.environ.get(variable_name)
+
+    if value is None:
+        return default
+
+    value = value.lower()
+
+    if value in ("true", "1"):
+        return True
+    elif value in ("false", "0"):
+        return False
+    else:
+        raise exceptions.InvalidValue(
+            'Environment variable "{}" must be one of "true", "false", "1", or "0".'.format(
+                variable_name
+            )
+        )
+
+
 def is_python_3():
     """Check if the Python interpreter is Python 2 or 3.
 

google/auth/credentials.py

@@ -368,16 +368,15 @@ def _lookup_trust_boundary(self, request):
         from google.oauth2 import _client
 
         # Verify the trust boundary feature flag is enabled.
-        if (
-            os.getenv(environment_vars.GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED, "").lower()
-            != "true"
+        if not _helpers.get_bool_from_env(
+            environment_vars.GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED, default=False
         ):
             # Skip the lookup and return early if it's not explicitly enabled.
-            return
+            return None
 
         # Skip trust boundary flow for non-gdu universe domain.
         if self.universe_domain != DEFAULT_UNIVERSE_DOMAIN:
-            return
+            return None
 
         url = self._build_trust_boundary_lookup_url()
         if not url:

google/auth/impersonated_credentials.py

@@ -439,7 +439,6 @@ def _make_copy(self):
 
     @_helpers.copy_docstring(credentials.CredentialsWithTrustBoundary)
     def with_trust_boundary(self, trust_boundary):
-        """Returns a copy of these credentials with a modified trust boundary."""
         cred = self._make_copy()
         cred._trust_boundary = trust_boundary
         return cred

google/oauth2/service_account.py

@@ -388,7 +388,6 @@ def with_token_uri(self, token_uri):
 
     @_helpers.copy_docstring(credentials.CredentialsWithTrustBoundary)
     def with_trust_boundary(self, trust_boundary):
-        """Returns a copy of these credentials with a modified trust boundary."""
         cred = self._make_copy()
         cred._trust_boundary = trust_boundary
         return cred

tests/oauth2/test_service_account.py

@@ -763,6 +763,52 @@ def test_refresh_trust_boundary_lookup_fails_no_cache(
         assert credentials._trust_boundary is None
         mock_lookup_trust_boundary.assert_called_once()
 
+    @mock.patch("google.oauth2._client.lookup_trust_boundary")
+    @mock.patch("google.oauth2._client.jwt_grant", autospec=True)
+    def test_refresh_trust_boundary_lookup_fails_with_cached_data(
+        self, mock_jwt_grant, mock_lookup_trust_boundary
+    ):
+        # Initial setup: Credentials with no trust boundary.
+        credentials = self.make_credentials(trust_boundary=None)
+        token = "token"
+        mock_jwt_grant.return_value = (
+            token,
+            _helpers.utcnow() + datetime.timedelta(seconds=500),
+            {},
+        )
+        request = mock.create_autospec(transport.Request, instance=True)
+
+        # First refresh: Successfully fetch a valid trust boundary.
+        mock_lookup_trust_boundary.return_value = self.VALID_TRUST_BOUNDARY
+        with mock.patch.dict(
+            os.environ, {environment_vars.GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED: "true"}
+        ):
+            credentials.refresh(request)
+
+        assert credentials.valid
+        assert credentials.token == token
+        assert credentials._trust_boundary == self.VALID_TRUST_BOUNDARY
+        mock_lookup_trust_boundary.assert_called_once_with(
+            request, self.EXPECTED_TRUST_BOUNDARY_LOOKUP_URL_DEFAULT_UNIVERSE, mock.ANY
+        )
+
+        # Second refresh: Mock lookup to fail, but expect cached data to be preserved.
+        mock_lookup_trust_boundary.reset_mock()
+        mock_lookup_trust_boundary.side_effect = exceptions.RefreshError(
+            "Lookup failed"
+        )
+
+        with mock.patch.dict(
+            os.environ, {environment_vars.GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED: "true"}
+        ):
+            credentials.refresh(request)  # This should NOT raise an exception
+
+        assert credentials.valid  # Credentials should still be valid
+        assert (
+            credentials._trust_boundary == self.VALID_TRUST_BOUNDARY
+        )  # Cached data should be preserved
+        mock_lookup_trust_boundary.assert_called_once()  # Lookup should have been attempted again
+
 
 class TestIDTokenCredentials(object):
     SERVICE_ACCOUNT_EMAIL = "service-account@example.com"

tests/test_credentials.py

@@ -16,7 +16,8 @@
 
 import mock
 import pytest  # type: ignore
-from google.auth import _client, exceptions
+from google.auth import exceptions
+from google.oauth2 import _client
 
 from google.auth import _helpers
 from google.auth import credentials
@@ -430,7 +431,6 @@ def test_lookup_trust_boundary_build_url_returns_none(self, mock_lookup_tb):
         with mock.patch.dict(
             os.environ, {environment_vars.GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED: "true"}
         ):
-            result = creds._lookup_trust_boundary(request)
             with pytest.raises(
                 exceptions.InvalidValue,
                 match="Failed to build trust boundary lookup URL.",

tests/test_impersonated_credentials.py

@@ -436,6 +436,54 @@ def test_refresh_starts_with_no_op_trust_boundary_skips_lookup(
         credentials.apply(headers_applied)
         assert headers_applied["x-allowed-locations"] == ""
 
+    @mock.patch("google.oauth2._client.lookup_trust_boundary")
+    def test_refresh_trust_boundary_lookup_fails_with_cached_data2(
+        self, mock_lookup_trust_boundary, mock_donor_credentials
+    ):
+        # Start with no trust boundary
+        credentials = self.make_credentials(lifetime=None, trust_boundary=None)
+        token = "token"
+
+        expire_time = (
+            _helpers.utcnow().replace(microsecond=0) + datetime.timedelta(seconds=500)
+        ).isoformat("T") + "Z"
+        response_body = {"accessToken": token, "expireTime": expire_time}
+
+        request = self.make_request(
+            data=json.dumps(response_body),
+            status=http_client.OK,
+        )
+
+        # First refresh: Successfully fetch a valid trust boundary.
+        mock_lookup_trust_boundary.return_value = self.VALID_TRUST_BOUNDARY
+        with mock.patch.dict(
+            os.environ, {environment_vars.GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED: "true"}
+        ), mock.patch(
+            "google.auth.metrics.token_request_access_token_impersonate",
+            return_value=ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE,
+        ):
+            credentials.refresh(request)
+
+        assert credentials.valid
+        # Verify trust boundary was set.
+        assert credentials._trust_boundary == self.VALID_TRUST_BOUNDARY
+        mock_lookup_trust_boundary.assert_called_once()
+
+        # Second refresh: Mock lookup to fail, but expect cached data to be preserved.
+        mock_lookup_trust_boundary.reset_mock()
+        mock_lookup_trust_boundary.side_effect = exceptions.RefreshError(
+            "Lookup failed"
+        )
+
+        with mock.patch.dict(
+            os.environ, {environment_vars.GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED: "true"}
+        ):
+            credentials.refresh(request)
+
+        assert credentials.valid
+        assert credentials._trust_boundary == self.VALID_TRUST_BOUNDARY
+        mock_lookup_trust_boundary.assert_called_once()
+
     @pytest.mark.parametrize("use_data_bytes", [True, False])
     def test_refresh_with_subject_success(self, use_data_bytes, mock_dwd_credentials):
         credentials = self.make_credentials(subject="test@email.com", lifetime=None)