Additional unit tests and update some old ones.

Author: nbayati

google/oauth2/service_account.py

@@ -482,10 +482,14 @@ def refresh(self, request):
             google.auth.exceptions.RefreshError: If the credentials could
                 not be refreshed.
         """
-        # Determine if we're going to use a self-signed JWT.
-        use_ssjwt = self._use_self_signed_jwt()
-
-        if use_ssjwt and self._is_trust_boundary_lookup_required():
+        # This is a special path for self-signed JWTs that need to look up a trust boundary.
+        # The `_subject` check is to ensure we are not in a domain-wide
+        # delegation flow, which uses a different authentication mechanism.
+        if (
+            self._always_use_jwt_access
+            and self._subject is None
+            and self._is_trust_boundary_lookup_required()
+        ):
             # Special case: self-signed JWT with trust boundary.
             # 1. Create a temporary self-signed JWT for the IAM API.
             iam_audience = "https://iamcredentials.{}/".format(self._universe_domain)
@@ -496,9 +500,12 @@ def refresh(self, request):
             # We temporarily set self.token for the base lookup method.
             # The base lookup method will call self.apply() which adds the
             # authorization header.
-            with _helpers.update_property(self, "token", iam_jwt_creds.token.decode()):
-                # This will call _lookup_trust_boundary and set self._trust_boundary
+            original_token = self.token
+            self.token = iam_jwt_creds.token.decode()
+            try:
                 self._refresh_trust_boundary(request)
+            finally:
+                self.token = original_token
 
             # 3. Now, refresh the original self-signed JWT for the target API.
             self._refresh_token(request)

tests/compute_engine/test_credentials.py

@@ -216,6 +216,18 @@ def test_with_universe_domain(self):
         assert creds.universe_domain == "universe_domain"
         assert creds._universe_domain_cached
 
+    def test_with_trust_boundary(self):
+        creds = self.credentials_with_all_fields
+        new_boundary = {"encodedLocations": "new_boundary"}
+        new_creds = creds.with_trust_boundary(new_boundary)
+
+        assert new_creds is not creds
+        assert new_creds._trust_boundary == new_boundary
+        assert new_creds._service_account_email == creds._service_account_email
+        assert new_creds._quota_project_id == creds._quota_project_id
+        assert new_creds._scopes == creds._scopes
+        assert new_creds._default_scopes == creds._default_scopes
+
     def test_token_usage_metrics(self):
         self.credentials.token = "token"
         self.credentials.expiry = None

tests/oauth2/test_service_account.py

@@ -270,6 +270,18 @@ def test__with_always_use_jwt_access_non_default_universe_domain(self):
             "always_use_jwt_access should be True for non-default universe domain"
         )
 
+    def test_with_trust_boundary(self):
+        credentials = self.make_credentials()
+        new_boundary = {"encodedLocations": "new_boundary"}
+        new_credentials = credentials.with_trust_boundary(new_boundary)
+
+        assert new_credentials is not credentials
+        assert new_credentials._trust_boundary == new_boundary
+        assert new_credentials._signer == credentials._signer
+        assert (
+            new_credentials.service_account_email == credentials.service_account_email
+        )
+
     def test__make_authorization_grant_assertion(self):
         credentials = self.make_credentials()
         token = credentials._make_authorization_grant_assertion()
@@ -851,6 +863,74 @@ def test_refresh_trust_boundary_lookup_fails_with_cached_data(
             },
         )  # Lookup should have been attempted again
 
+    @mock.patch("google.oauth2._client._lookup_trust_boundary")
+    def test_refresh_with_self_signed_jwt_and_trust_boundary(
+        self, mock_lookup_trust_boundary
+    ):
+        # --- Setup ---
+        # Create credentials that use self-signed JWT and have no initial boundary.
+        credentials = self.make_credentials(trust_boundary=None)
+        credentials._always_use_jwt_access = True
+        credentials._scopes = ["https://www.googleapis.com/auth/devstorage.read_only"]
+        request = mock.create_autospec(transport.Request, instance=True)
+
+        # Mock the trust boundary lookup to return a valid value.
+        mock_lookup_trust_boundary.return_value = self.VALID_TRUST_BOUNDARY
+
+        # Mock the two JWTs that will be created: one for the IAM lookup
+        # and one for the final token.
+        iam_jwt_creds_mock = mock.MagicMock(spec=jwt.Credentials, token=b"iam_token")
+        final_jwt_creds_mock = mock.MagicMock(
+            spec=jwt.Credentials,
+            token=b"final_token",
+            expiry=_helpers.utcnow() + datetime.timedelta(hours=1),
+        )
+        from_signing_credentials_mock = mock.patch(
+            "google.auth.jwt.Credentials.from_signing_credentials",
+            side_effect=[iam_jwt_creds_mock, final_jwt_creds_mock],
+        )
+        env_mock = mock.patch.dict(
+            os.environ, {environment_vars.GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED: "true"}
+        )
+
+        # Refresh credential to trigger creating the self-signed jwt token and
+        # fetching the trust boundary data.
+        with from_signing_credentials_mock as mock_from_signing, env_mock:
+            credentials.refresh(request)
+
+        # --- Assert ---
+        # 1. Verify the temporary IAM JWT was created and used for the lookup.
+        iam_call = mock_from_signing.call_args_list[0]
+        assert iam_call.args[0] is credentials
+        assert iam_call.args[1] == "https://iamcredentials.googleapis.com/"
+        iam_jwt_creds_mock.refresh.assert_called_once_with(request)
+
+        mock_lookup_trust_boundary.assert_called_once()
+        _, lookup_kwargs = mock_lookup_trust_boundary.call_args
+        assert lookup_kwargs["headers"]["authorization"] == "Bearer iam_token"
+
+        # 2. Verify the final, audience-specific JWT was created and refreshed.
+        final_token_call = mock_from_signing.call_args_list[1]
+        assert final_token_call.args[0] is credentials
+        assert final_token_call.args[1] is None  # Audience is derived from scopes
+        assert final_token_call.kwargs["additional_claims"] == {
+            "scope": " ".join(credentials._scopes)
+        }
+        final_jwt_creds_mock.refresh.assert_called_once_with(request)
+
+        # 3. Verify the final state of the credentials object.
+        assert credentials.token == "final_token"
+        assert credentials._trust_boundary == self.VALID_TRUST_BOUNDARY
+
+    def test_build_trust_boundary_lookup_url_no_email(self):
+        credentials = self.make_credentials()
+        credentials._service_account_email = None
+
+        with pytest.raises(ValueError) as excinfo:
+            credentials._build_trust_boundary_lookup_url()
+
+        assert "Service account email is required" in str(excinfo.value)
+
 
 class TestIDTokenCredentials(object):
     SERVICE_ACCOUNT_EMAIL = "service-account@example.com"
@@ -981,6 +1061,19 @@ def test_with_token_uri(self):
         creds_with_new_token_uri = credentials.with_token_uri(new_token_uri)
         assert creds_with_new_token_uri._token_uri == new_token_uri
 
+    def test_with_trust_boundary(self):
+        credentials = self.make_credentials()
+        new_boundary = {"encodedLocations": "new_boundary"}
+        new_credentials = credentials.with_trust_boundary(new_boundary)
+
+        assert new_credentials is not credentials
+        assert new_credentials._trust_boundary == new_boundary
+        assert new_credentials._signer == credentials._signer
+        assert (
+            new_credentials.service_account_email == credentials.service_account_email
+        )
+        assert new_credentials._target_audience == credentials._target_audience
+
     def test__make_authorization_grant_assertion(self):
         credentials = self.make_credentials()
         token = credentials._make_authorization_grant_assertion()

tests/test_aws.py

@@ -42,8 +42,10 @@
 SERVICE_ACCOUNT_IMPERSONATION_URL_BASE = (
     "https://us-east1-iamcredentials.googleapis.com"
 )
-SERVICE_ACCOUNT_IMPERSONATION_URL_ROUTE = "/v1/projects/-/serviceAccounts/{}:generateAccessToken".format(
-    SERVICE_ACCOUNT_EMAIL
+SERVICE_ACCOUNT_IMPERSONATION_URL_ROUTE = (
+    "/v1/projects/-/serviceAccounts/{}:generateAccessToken".format(
+        SERVICE_ACCOUNT_EMAIL
+    )
 )
 SERVICE_ACCOUNT_IMPERSONATION_URL = (
     SERVICE_ACCOUNT_IMPERSONATION_URL_BASE + SERVICE_ACCOUNT_IMPERSONATION_URL_ROUTE
@@ -920,7 +922,7 @@ def assert_token_request_kwargs(
         assert request_kwargs["body"] is not None
         body_tuples = urllib.parse.parse_qsl(request_kwargs["body"])
         assert len(body_tuples) == len(request_data.keys())
-        for (k, v) in body_tuples:
+        for k, v in body_tuples:
             assert v.decode("utf-8") == request_data[k.decode("utf-8")]
 
     @classmethod
@@ -1344,9 +1346,9 @@ def test_retrieve_subject_token_success_temp_creds_no_environment_vars_idmsv2(
             imdsv2_session_token_data=self.AWS_IMDSV2_SESSION_TOKEN,
         )
         credential_source_token_url = self.CREDENTIAL_SOURCE.copy()
-        credential_source_token_url[
-            "imdsv2_session_token_url"
-        ] = IMDSV2_SESSION_TOKEN_URL
+        credential_source_token_url["imdsv2_session_token_url"] = (
+            IMDSV2_SESSION_TOKEN_URL
+        )
         credentials = self.make_credentials(
             credential_source=credential_source_token_url
         )
@@ -1452,9 +1454,9 @@ def test_retrieve_subject_token_success_temp_creds_environment_vars_missing_secr
             imdsv2_session_token_data=self.AWS_IMDSV2_SESSION_TOKEN,
         )
         credential_source_token_url = self.CREDENTIAL_SOURCE.copy()
-        credential_source_token_url[
-            "imdsv2_session_token_url"
-        ] = IMDSV2_SESSION_TOKEN_URL
+        credential_source_token_url["imdsv2_session_token_url"] = (
+            IMDSV2_SESSION_TOKEN_URL
+        )
         credentials = self.make_credentials(
             credential_source=credential_source_token_url
         )
@@ -1509,9 +1511,9 @@ def test_retrieve_subject_token_success_temp_creds_environment_vars_missing_acce
             imdsv2_session_token_data=self.AWS_IMDSV2_SESSION_TOKEN,
         )
         credential_source_token_url = self.CREDENTIAL_SOURCE.copy()
-        credential_source_token_url[
-            "imdsv2_session_token_url"
-        ] = IMDSV2_SESSION_TOKEN_URL
+        credential_source_token_url["imdsv2_session_token_url"] = (
+            IMDSV2_SESSION_TOKEN_URL
+        )
         credentials = self.make_credentials(
             credential_source=credential_source_token_url
         )
@@ -1560,9 +1562,9 @@ def test_retrieve_subject_token_success_temp_creds_environment_vars_missing_cred
             imdsv2_session_token_data=self.AWS_IMDSV2_SESSION_TOKEN,
         )
         credential_source_token_url = self.CREDENTIAL_SOURCE.copy()
-        credential_source_token_url[
-            "imdsv2_session_token_url"
-        ] = IMDSV2_SESSION_TOKEN_URL
+        credential_source_token_url["imdsv2_session_token_url"] = (
+            IMDSV2_SESSION_TOKEN_URL
+        )
         credentials = self.make_credentials(
             credential_source=credential_source_token_url
         )
@@ -1611,9 +1613,9 @@ def test_retrieve_subject_token_success_temp_creds_idmsv2(self, utcnow):
             role_status=http_client.OK, role_name=self.AWS_ROLE
         )
         credential_source_token_url = self.CREDENTIAL_SOURCE.copy()
-        credential_source_token_url[
-            "imdsv2_session_token_url"
-        ] = IMDSV2_SESSION_TOKEN_URL
+        credential_source_token_url["imdsv2_session_token_url"] = (
+            IMDSV2_SESSION_TOKEN_URL
+        )
         credentials = self.make_credentials(
             credential_source=credential_source_token_url
         )
@@ -1692,9 +1694,9 @@ def test_retrieve_subject_token_session_error_idmsv2(self, utcnow):
             imdsv2_session_token_data="unauthorized",
         )
         credential_source_token_url = self.CREDENTIAL_SOURCE.copy()
-        credential_source_token_url[
-            "imdsv2_session_token_url"
-        ] = IMDSV2_SESSION_TOKEN_URL
+        credential_source_token_url["imdsv2_session_token_url"] = (
+            IMDSV2_SESSION_TOKEN_URL
+        )
         credentials = self.make_credentials(
             credential_source=credential_source_token_url
         )
@@ -2057,7 +2059,9 @@ def test_refresh_success_with_impersonation_ignore_default_scopes(
             "authorization": "Bearer {}".format(self.SUCCESS_RESPONSE["access_token"]),
             "x-goog-user-project": QUOTA_PROJECT_ID,
             "x-goog-api-client": IMPERSONATE_ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE,
-            "x-allowed-locations": "0x0",
+            # TODO(negarb): Uncomment and update when trust boundary is supported
+            # for external account credentials.
+            # "x-allowed-locations": "0x0",
         }
         impersonation_request_data = {
             "delegates": None,
@@ -2150,7 +2154,7 @@ def test_refresh_success_with_impersonation_use_default_scopes(
             "authorization": "Bearer {}".format(self.SUCCESS_RESPONSE["access_token"]),
             "x-goog-user-project": QUOTA_PROJECT_ID,
             "x-goog-api-client": IMPERSONATE_ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE,
-            "x-allowed-locations": "0x0",
+            # "x-allowed-locations": "0x0",
         }
         impersonation_request_data = {
             "delegates": None,
@@ -2345,7 +2349,7 @@ def test_refresh_success_with_supplier_with_impersonation(
             "authorization": "Bearer {}".format(self.SUCCESS_RESPONSE["access_token"]),
             "x-goog-user-project": QUOTA_PROJECT_ID,
             "x-goog-api-client": IMPERSONATE_ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE,
-            "x-allowed-locations": "0x0",
+            # "x-allowed-locations": "0x0",
         }
         impersonation_request_data = {
             "delegates": None,

tests/test_credentials.py

@@ -26,7 +26,7 @@
 
 
 class CredentialsImpl(credentials.CredentialsWithTrustBoundary):
-    def refresh(self, request):
+    def _refresh_token(self, request):
         self.token = request
         self.expiry = (
             datetime.datetime.utcnow()
@@ -354,7 +354,7 @@ def test_token_state_no_expiry():
 
 
 class TestCredentialsWithTrustBoundary(object):
-    @mock.patch.object(_client, "lookup_trust_boundary")
+    @mock.patch.object(_client, "_lookup_trust_boundary")
     def test_lookup_trust_boundary_env_var_not_true(self, mock_lookup_tb):
         creds = CredentialsImpl()
         request = mock.Mock()
@@ -363,12 +363,12 @@ def test_lookup_trust_boundary_env_var_not_true(self, mock_lookup_tb):
         with mock.patch.dict(
             os.environ, {environment_vars.GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED: "false"}
         ):
-            result = creds._lookup_trust_boundary(request)
+            result = creds._refresh_trust_boundary(request)
 
         assert result is None
         mock_lookup_tb.assert_not_called()
 
-    @mock.patch.object(_client, "lookup_trust_boundary")
+    @mock.patch.object(_client, "_lookup_trust_boundary")
     def test_lookup_trust_boundary_env_var_missing(self, mock_lookup_tb):
         creds = CredentialsImpl()
         request = mock.Mock()
@@ -378,12 +378,12 @@ def test_lookup_trust_boundary_env_var_missing(self, mock_lookup_tb):
             # Remove the var if it was set by other tests
             if environment_vars.GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED in os.environ:
                 del os.environ[environment_vars.GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED]
-            result = creds._lookup_trust_boundary(request)
+            result = creds._refresh_trust_boundary(request)
 
         assert result is None
         mock_lookup_tb.assert_not_called()
 
-    @mock.patch.object(_client, "lookup_trust_boundary")
+    @mock.patch.object(_client, "_lookup_trust_boundary")
     def test_lookup_trust_boundary_non_default_universe(self, mock_lookup_tb):
         creds = CredentialsImpl()
         creds._universe_domain = "my.universe.com"  # Non-GDU
@@ -392,12 +392,12 @@ def test_lookup_trust_boundary_non_default_universe(self, mock_lookup_tb):
         with mock.patch.dict(
             os.environ, {environment_vars.GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED: "true"}
         ):
-            result = creds._lookup_trust_boundary(request)
+            result = creds._refresh_trust_boundary(request)
 
         assert result is None
         mock_lookup_tb.assert_not_called()
 
-    @mock.patch.object(_client, "lookup_trust_boundary")
+    @mock.patch.object(_client, "_lookup_trust_boundary")
     def test_lookup_trust_boundary_calls_client_and_build_url(self, mock_lookup_tb):
         creds = CredentialsImpl()
         creds.token = "test_token"  # For _build_trust_boundary_lookup_url
@@ -422,7 +422,7 @@ def test_lookup_trust_boundary_calls_client_and_build_url(self, mock_lookup_tb):
             request, expected_url, headers=expected_headers
         )
 
-    @mock.patch.object(_client, "lookup_trust_boundary")
+    @mock.patch.object(_client, "_lookup_trust_boundary")
     def test_lookup_trust_boundary_build_url_returns_none(self, mock_lookup_tb):
         creds = CredentialsImpl()
         request = mock.Mock()