diff --git a/fix_all.py b/fix_all.py new file mode 100644 index 000000000000..10d61256d5d7 --- /dev/null +++ b/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 + diff --git a/fix_oauth_cred.py b/fix_oauth_cred.py new file mode 100644 index 000000000000..b3065bba6e57 --- /dev/null +++ b/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") diff --git a/fix_service_acct.py b/fix_service_acct.py new file mode 100644 index 000000000000..a5d3e9548030 --- /dev/null +++ b/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") diff --git a/fix_test.py b/fix_test.py new file mode 100644 index 000000000000..8cab1abb0ab5 --- /dev/null +++ b/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 diff --git a/packages/google-auth/tests/compute_engine/test_credentials.py b/packages/google-auth/tests/compute_engine/test_credentials.py index 7fb2b8b504fc..4cc0ce25e46d 100644 --- a/packages/google-auth/tests/compute_engine/test_credentials.py +++ b/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), diff --git a/packages/google-auth/tests/oauth2/test_credentials.py b/packages/google-auth/tests/oauth2/test_credentials.py index 5a1ec2f757ad..43df9b3a0cc4 100644 --- a/packages/google-auth/tests/oauth2/test_credentials.py +++ b/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") diff --git a/packages/google-auth/tests/oauth2/test_service_account.py b/packages/google-auth/tests/oauth2/test_service_account.py index 958eace2dd22..1d70543057d7 100644 --- a/packages/google-auth/tests/oauth2/test_service_account.py +++ b/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() diff --git a/packages/google-auth/tests/test_external_account.py b/packages/google-auth/tests/test_external_account.py index 870b07d47b6e..d0f90cb940ed 100644 --- a/packages/google-auth/tests/test_external_account.py +++ b/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( diff --git a/packages/google-auth/tests_async/oauth2/test_credentials_async.py b/packages/google-auth/tests_async/oauth2/test_credentials_async.py index 0a5d8ab1aaf9..d8bf82a0b59f 100644 --- a/packages/google-auth/tests_async/oauth2/test_credentials_async.py +++ b/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() diff --git a/packages/google-auth/tests_async/oauth2/test_service_account_async.py b/packages/google-auth/tests_async/oauth2/test_service_account_async.py index e0c2e0d60a60..0539ecc80e13 100644 --- a/packages/google-auth/tests_async/oauth2/test_service_account_async.py +++ b/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() diff --git a/patch_tests.py b/patch_tests.py new file mode 100644 index 000000000000..3d41494ead60 --- /dev/null +++ b/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) diff --git a/print_tests.py b/print_tests.py new file mode 100644 index 000000000000..cb125d792a34 --- /dev/null +++ b/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]) diff --git a/test_script.py b/test_script.py new file mode 100644 index 000000000000..17fd06dfbe97 --- /dev/null +++ b/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)