Refactor custom headers implementation and tests

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",
+        metrics.API_CLIENT_HEADER,
+        "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"""

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"