test(auth): assert quota project header injection in google-auth tests

fix_all.py

@@ -0,0 +1,29 @@
+import os
+import re
+
+TEST_FILES = [
+    "packages/google-auth/tests/compute_engine/test_credentials.py",
+    "packages/google-auth/tests/oauth2/test_credentials.py",
+    "packages/google-auth/tests/oauth2/test_service_account.py",
+    "packages/google-auth/tests/test_app_engine.py",
+    "packages/google-auth/tests/test_downscoped.py",
+    "packages/google-auth/tests/test_external_account.py",
+    "packages/google-auth/tests/test_external_account_authorized_user.py",
+    "packages/google-auth/tests/test_impersonated_credentials.py",
+    "packages/google-auth/tests/test_jwt.py",
+    "packages/google-auth/tests_async/oauth2/test_credentials_async.py",
+    "packages/google-auth/tests_async/oauth2/test_service_account_async.py",
+    "packages/google-auth/tests_async/test_jwt_async.py",
+]
+
+def process_file(filepath):
+    with open(filepath, 'r') as f:
+        content = f.read()
+
+    # We want to find `test_with_quota_project` and insert header assertions.
+    # But since each test names its variables differently (`creds`, `new_creds`, `quota_project_creds`, `self.credentials`),
+    # it's tricky to automate completely with regex. Let's do a semi-manual replacement for the most common ones.
+
+if __name__ == "__main__":
+    pass
+

fix_oauth_cred.py

@@ -0,0 +1,26 @@
+import re
+
+path = "packages/google-auth/tests/oauth2/test_credentials.py"
+with open(path, "r") as f:
+    content = f.read()
+
+# Replace creds.apply(headers) with new_creds.before_request(...)
+old = """        new_creds = creds.with_quota_project("new-project-456")
+        assert new_creds.quota_project_id == "new-project-456"
+        headers = {}
+        creds.apply(headers)
+        assert "x-goog-user-project" in headers"""
+
+new = """        new_creds = creds.with_quota_project("new-project-456")
+        assert new_creds.quota_project_id == "new-project-456"
+
+        request = mock.create_autospec(transport.Request, instance=True)
+        headers = {}
+        new_creds.before_request(request, "GET", "https://example.com", headers)
+        assert headers.get("x-goog-user-project") == "new-project-456" """
+
+if old in content:
+    content = content.replace(old, new)
+    with open(path, "w") as f:
+        f.write(content)
+    print("Replaced in oauth2/test_credentials.py")

fix_service_acct.py

@@ -0,0 +1,23 @@
+path = "packages/google-auth/tests/oauth2/test_service_account.py"
+with open(path, "r") as f:
+    content = f.read()
+
+old = """        new_credentials = credentials.with_quota_project("new-project-456")
+        assert new_credentials.quota_project_id == "new-project-456"
+        hdrs = {}
+        new_credentials.apply(hdrs, token="tok")
+        assert "x-goog-user-project" in hdrs"""
+
+new = """        new_credentials = credentials.with_quota_project("new-project-456")
+        assert new_credentials.quota_project_id == "new-project-456"
+
+        request = mock.create_autospec(transport.Request, instance=True)
+        hdrs = {}
+        new_credentials.before_request(request, "GET", "https://example.com", hdrs)
+        assert hdrs.get("x-goog-user-project") == "new-project-456" """
+
+if old in content:
+    content = content.replace(old, new)
+    with open(path, "w") as f:
+        f.write(content)
+    print("Replaced in oauth2/test_service_account.py")

fix_test.py

@@ -0,0 +1,28 @@
+import re
+
+files = [
+    "packages/google-auth/tests/compute_engine/test_credentials.py",
+    "packages/google-auth/tests/oauth2/test_credentials.py",
+    "packages/google-auth/tests/oauth2/test_service_account.py",
+    "packages/google-auth/tests/test_app_engine.py",
+    "packages/google-auth/tests/test_downscoped.py",
+    "packages/google-auth/tests/test_external_account.py",
+    "packages/google-auth/tests/test_external_account_authorized_user.py",
+    "packages/google-auth/tests/test_impersonated_credentials.py",
+    "packages/google-auth/tests/test_jwt.py",
+    "packages/google-auth/tests_async/oauth2/test_credentials_async.py",
+    "packages/google-auth/tests_async/oauth2/test_service_account_async.py",
+    "packages/google-auth/tests_async/test_jwt_async.py",
+]
+
+for file in files:
+    with open(file, 'r') as f:
+        content = f.read()
+
+    # Pattern 1: creds.apply(headers) -> we should add assertion
+    content = re.sub(
+        r'(\s+)creds\.apply\(headers\)\s+assert "x-goog-user-project" in headers',
+        r'\1creds.apply(headers)\1assert headers.get("x-goog-user-project") == "new-project-456" if "new-project-456" in content else headers.get("x-goog-user-project") == "quota-project-123"',
+        content
+    )
+    # Actually it's better to just write a simple script that targets exactly what we need

packages/google-auth/tests/compute_engine/test_credentials.py

@@ -860,6 +860,11 @@ def test_with_quota_project(self, sign, get, utcnow):
         # Check that the signer have been initialized with a Request object
         assert isinstance(self.credentials._signer._request, transport.Request)
 
+        headers = {}
+        self.credentials.token = "fake-token"
+        self.credentials.before_request(request, "GET", "https://example.com", headers)
+        assert headers.get("x-goog-user-project") == "project-foo"
+
     @mock.patch(
         "google.auth._helpers.utcnow",
         return_value=_helpers.utcfromtimestamp(0),

packages/google-auth/tests/oauth2/test_credentials.py

@@ -844,9 +844,10 @@ def test_with_quota_project(self):
 
         new_creds = creds.with_quota_project("new-project-456")
         assert new_creds.quota_project_id == "new-project-456"
+        request = mock.create_autospec(transport.Request, instance=True)
         headers = {}
-        creds.apply(headers)
-        assert "x-goog-user-project" in headers
+        new_creds.before_request(request, "GET", "https://example.com", headers)
+        assert headers.get("x-goog-user-project") == "new-project-456"
 
     def test_with_universe_domain(self):
         creds = credentials.Credentials(token="token")

packages/google-auth/tests/oauth2/test_service_account.py

@@ -224,9 +224,11 @@ def test_with_quota_project(self):
         credentials = self.make_credentials()
         new_credentials = credentials.with_quota_project("new-project-456")
         assert new_credentials.quota_project_id == "new-project-456"
+        request = mock.create_autospec(transport.Request, instance=True)
         hdrs = {}
-        new_credentials.apply(hdrs, token="tok")
-        assert "x-goog-user-project" in hdrs
+        new_credentials.token = "tok"
+        new_credentials.before_request(request, "GET", "https://example.com", hdrs)
+        assert hdrs.get("x-goog-user-project") == "new-project-456"
 
     def test_copy_regional_access_boundary_manager_state_and_config_with_scopes(self):
         credentials = self.make_credentials()

packages/google-auth/tests/test_external_account.py

@@ -456,6 +456,13 @@ def test_with_quota_project(self):
         quota_project_creds = credentials.with_quota_project("project-foo")
 
         assert quota_project_creds.quota_project_id == "project-foo"
+        request = mock.create_autospec(transport.Request, instance=True)
+        headers = {}
+        quota_project_creds.token = "fake-token"
+        quota_project_creds.before_request(
+            request, "GET", "https://example.com", headers
+        )
+        assert headers.get("x-goog-user-project") == "project-foo"
 
     def test_with_quota_project_workforce_pool(self):
         credentials = self.make_workforce_pool_credentials(

packages/google-auth/tests_async/oauth2/test_credentials_async.py

@@ -23,6 +23,7 @@
 
 from google.auth import _helpers
 from google.auth import exceptions
+from google.auth import transport
 from google.oauth2 import _credentials_async as _credentials_async
 from google.oauth2 import credentials
 from tests.oauth2 import test_credentials
@@ -344,7 +345,8 @@ def test_apply_with_no_quota_project_id(self):
         creds.apply(headers)
         assert "x-goog-user-project" not in headers
 
-    def test_with_quota_project(self):
+    @pytest.mark.asyncio
+    async def test_with_quota_project(self):
         creds = _credentials_async.Credentials(
             token="token",
             refresh_token=self.REFRESH_TOKEN,
@@ -356,9 +358,10 @@ def test_with_quota_project(self):
 
         new_creds = creds.with_quota_project("new-project-456")
         assert new_creds.quota_project_id == "new-project-456"
+        request = mock.create_autospec(transport.Request, instance=True)
         headers = {}
-        creds.apply(headers)
-        assert "x-goog-user-project" in headers
+        await new_creds.before_request(request, "GET", "https://example.com", headers)
+        assert headers.get("x-goog-user-project") == "new-project-456"
 
     def test_from_authorized_user_info(self):
         info = test_credentials.AUTH_USER_INFO.copy()

packages/google-auth/tests_async/oauth2/test_service_account_async.py

@@ -139,13 +139,18 @@ def test_with_claims(self):
         new_credentials = credentials.with_claims({"meep": "moop"})
         assert new_credentials._additional_claims == {"meep": "moop"}
 
-    def test_with_quota_project(self):
+    @pytest.mark.asyncio
+    async def test_with_quota_project(self):
         credentials = self.make_credentials()
         new_credentials = credentials.with_quota_project("new-project-456")
         assert new_credentials.quota_project_id == "new-project-456"
+        request = mock.create_autospec(transport.Request, instance=True)
         hdrs = {}
-        new_credentials.apply(hdrs, token="tok")
-        assert "x-goog-user-project" in hdrs
+        new_credentials.token = "tok"
+        await new_credentials.before_request(
+            request, "GET", "https://example.com", hdrs
+        )
+        assert hdrs.get("x-goog-user-project") == "new-project-456"
 
     def test__make_authorization_grant_assertion(self):
         credentials = self.make_credentials()

patch_tests.py

@@ -0,0 +1,18 @@
+import re
+import glob
+import os
+
+files = glob.glob("packages/google-auth/tests/**/*.py", recursive=True)
+
+for file in files:
+    with open(file, 'r') as f:
+        content = f.read()
+
+    # If it is oauth2/test_credentials.py
+    if file.endswith("oauth2/test_credentials.py"):
+        content = content.replace('creds.apply(headers)\n        assert "x-goog-user-project" in headers', 
+                                  'new_creds.apply(headers)\n        assert headers["x-goog-user-project"] == "new-project-456"')
+
+    # We will write back content
+    with open(file, 'w') as f:
+        f.write(content)

print_tests.py

@@ -0,0 +1,17 @@
+import os, glob
+
+for root, _, files in os.walk("packages/google-auth/tests"):
+    for file in files:
+        if file.endswith(".py"):
+            path = os.path.join(root, file)
+            with open(path, "r") as f:
+                content = f.read()
+            if "def test_with_quota_project" in content:
+                print(f"\n--- {path} ---")
+                lines = content.split('\n')
+                for i, line in enumerate(lines):
+                    if "def test_with_quota_project" in line:
+                        for j in range(i, min(i+15, len(lines))):
+                            if "def test_" in lines[j] and j > i:
+                                break
+                            print(lines[j])

test_script.py

@@ -0,0 +1,14 @@
+from unittest import mock
+import google.auth.transport.requests as transport
+from google.auth.compute_engine import credentials
+
+creds = credentials.Credentials(
+    service_account_email="test@example.com",
+    quota_project_id="project-foo",
+)
+creds.token = "fake-token" # Inject token to avoid refresh
+
+request = mock.create_autospec(transport.Request, instance=True)
+headers = {}
+creds.before_request(request, "GET", "http://example.com", headers)
+print(headers)