From e8f801c5872e2c7f020555cd76dc62959447c583 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Tue, 30 Dec 2025 14:50:36 +0000 Subject: [PATCH 1/2] Refactor custom headers feature and tests - Implement `with_headers` in `Credentials` using shallow copy for safety. - Move `PROTECTED_HEADERS` to class attribute `_PROTECTED_HEADERS`. - Refactor `tests/test_credentials.py` to use `TestWithHeaders` class and `pytest.mark.parametrize` for better coverage and readability. - Fix logic to ensure `_custom_headers` are properly isolated between copies. --- google/auth/credentials.py | 34 +++++++++++++++++++++ tests/test_credentials.py | 62 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 96 insertions(+) diff --git a/google/auth/credentials.py b/google/auth/credentials.py index 82c73c3bf..07e84e396 100644 --- a/google/auth/credentials.py +++ b/google/auth/credentials.py @@ -16,6 +16,7 @@ """Interfaces for credentials.""" import abc +import copy from enum import Enum import os from typing import List @@ -69,6 +70,14 @@ def __init__(self): self._use_non_blocking_refresh = False self._refresh_worker = RefreshThreadManager() + self._custom_headers = {} + + _PROTECTED_HEADERS = { + "authorization", + "x-goog-user-project", + "x-goog-api-client", + "x-allowed-locations", + } @property def expired(self): @@ -185,6 +194,7 @@ def apply(self, headers, token=None): self._apply(headers, token) if self.quota_project_id: headers["x-goog-user-project"] = self.quota_project_id + headers.update(self._custom_headers) def _blocking_refresh(self, request): if not self.valid: @@ -233,6 +243,30 @@ def before_request(self, request, method, url, headers): def with_non_blocking_refresh(self): self._use_non_blocking_refresh = True + def with_headers(self, headers): + """Returns a copy of these credentials with additional custom headers. + + Args: + headers (Mapping[str, str]): The custom headers to add. + + Returns: + google.auth.credentials.Credentials: A new credentials instance. + + Raises: + ValueError: If a protected header is included in the input headers. + """ + for key in headers: + if key.lower() in self._PROTECTED_HEADERS: + raise ValueError( + f"Header '{key}' is protected and cannot be set with with_headers. " + "These headers are managed by the library." + ) + + new_creds = copy.copy(self) + new_creds._custom_headers = self._custom_headers.copy() + new_creds._custom_headers.update(headers) + return new_creds + class CredentialsWithQuotaProject(Credentials): """Abstract base for credentials supporting ``with_quota_project`` factory""" diff --git a/tests/test_credentials.py b/tests/test_credentials.py index 750c92af2..895cffdff 100644 --- a/tests/test_credentials.py +++ b/tests/test_credentials.py @@ -72,6 +72,68 @@ def test_with_non_blocking_refresh(): assert c._use_non_blocking_refresh +class TestWithHeaders: + def test_add_new_header(self): + credentials = CredentialsImpl() + request = mock.Mock() + + creds_with_header = credentials.with_headers({"X-Custom-Header": "value1"}) + headers = {} + creds_with_header.before_request(request, "http://example.com", "GET", headers) + + assert headers["X-Custom-Header"] == "value1" + assert "authorization" in headers + # Ensure it is a new instance + assert creds_with_header is not credentials + # Ensure original credentials are not modified + assert ( + not hasattr(credentials, "_custom_headers") or not credentials._custom_headers + ) + + def test_update_existing_header(self): + credentials = CredentialsImpl() + request = mock.Mock() + + creds_with_header = credentials.with_headers({"X-Custom-Header": "value1"}) + creds_updated = creds_with_header.with_headers({"X-Custom-Header": "value2"}) + headers = {} + creds_updated.before_request(request, "http://example.com", "GET", headers) + + assert headers["X-Custom-Header"] == "value2" + + def test_chaining_headers(self): + credentials = CredentialsImpl() + request = mock.Mock() + + creds_chained = credentials.with_headers({"X-Header-1": "v1"}).with_headers( + {"X-Header-2": "v2"} + ) + headers = {} + creds_chained.before_request(request, "http://example.com", "GET", headers) + + assert headers["X-Header-1"] == "v1" + assert headers["X-Header-2"] == "v2" + + @pytest.mark.parametrize( + "header_key", + ["Authorization", "X-Goog-User-Project", "authorization", "x-allowed-locations"], + ) + def test_protected_headers(self, header_key): + credentials = CredentialsImpl() + with pytest.raises(ValueError, match="is protected and cannot be set"): + credentials.with_headers({header_key: "value"}) + + def test_original_credentials_not_modified(self): + credentials = CredentialsImpl() + request = mock.Mock() + + credentials.with_headers({"X-Custom-Header": "value1"}) + headers = {} + credentials.before_request(request, "http://example.com", "GET", headers) + + assert "X-Custom-Header" not in headers + + def test_expired_and_valid(): credentials = CredentialsImpl() credentials.token = "token" From e086f8e69ac1175759c58e5fd867c350a31e84e4 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Tue, 30 Dec 2025 15:03:42 +0000 Subject: [PATCH 2/2] Refactor custom headers feature and tests - Implement `with_headers` in `Credentials` using shallow copy for safety. - Move `PROTECTED_HEADERS` to class attribute `_PROTECTED_HEADERS`. - Use `metrics.API_CLIENT_HEADER` constant in `_PROTECTED_HEADERS`. - Refactor `tests/test_credentials.py` to use `TestWithHeaders` class and `pytest.mark.parametrize` for better coverage and readability. - Fix logic to ensure `_custom_headers` are properly isolated between copies. --- google/auth/credentials.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google/auth/credentials.py b/google/auth/credentials.py index 07e84e396..f1918fdcf 100644 --- a/google/auth/credentials.py +++ b/google/auth/credentials.py @@ -75,7 +75,7 @@ def __init__(self): _PROTECTED_HEADERS = { "authorization", "x-goog-user-project", - "x-goog-api-client", + metrics.API_CLIENT_HEADER, "x-allowed-locations", }