From a143760fdb1ca1225a60b5c7901d5db99ab6cdd9 Mon Sep 17 00:00:00 2001 From: Sijun Liu Date: Tue, 7 Nov 2023 14:45:40 -0800 Subject: [PATCH] fix: add with_universe_domain to service account and external cred --- google/auth/external_account.py | 16 ++++++++++++++++ google/oauth2/service_account.py | 16 ++++++++++++++++ tests/oauth2/test_service_account.py | 11 +++++++++++ tests/test_external_account.py | 5 +++++ 4 files changed, 48 insertions(+) diff --git a/google/auth/external_account.py b/google/auth/external_account.py index 28b004c5f..e7fed8695 100644 --- a/google/auth/external_account.py +++ b/google/auth/external_account.py @@ -415,6 +415,22 @@ def with_token_uri(self, token_uri): new_cred._metrics_options = self._metrics_options return new_cred + def with_universe_domain(self, universe_domain): + """Create a copy of these credentials with the given universe domain. + + Args: + universe_domain (str): The universe domain value. + + Returns: + google.auth.external_account.Credentials: A new credentials + instance. + """ + kwargs = self._constructor_args() + kwargs.update(universe_domain=universe_domain) + new_cred = self.__class__(**kwargs) + new_cred._metrics_options = self._metrics_options + return new_cred + def _initialize_impersonated_credentials(self): """Generates an impersonated credentials. diff --git a/google/oauth2/service_account.py b/google/oauth2/service_account.py index 803b13070..fcdd547f7 100644 --- a/google/oauth2/service_account.py +++ b/google/oauth2/service_account.py @@ -328,6 +328,22 @@ def with_always_use_jwt_access(self, always_use_jwt_access): cred._always_use_jwt_access = always_use_jwt_access return cred + def with_universe_domain(self, universe_domain): + """Create a copy of these credentials with the given universe domain. + + Args: + universe_domain (str): The universe domain value. + + Returns: + google.auth.service_account.Credentials: A new credentials + instance. + """ + cred = self._make_copy() + cred._universe_domain = universe_domain + if universe_domain != _DEFAULT_UNIVERSE_DOMAIN: + cred._always_use_jwt_access = True + return cred + def with_subject(self, subject): """Create a copy of these credentials with the specified subject. diff --git a/tests/oauth2/test_service_account.py b/tests/oauth2/test_service_account.py index b963b157c..898259ce8 100644 --- a/tests/oauth2/test_service_account.py +++ b/tests/oauth2/test_service_account.py @@ -205,6 +205,17 @@ 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_universe_domain(self): + credentials = self.make_credentials() + + new_credentials = credentials.with_universe_domain("dummy_universe.com") + assert new_credentials.universe_domain == "dummy_universe.com" + assert new_credentials._always_use_jwt_access + + new_credentials = credentials.with_universe_domain("googleapis.com") + assert new_credentials.universe_domain == "googleapis.com" + assert not new_credentials._always_use_jwt_access + def test__with_always_use_jwt_access(self): credentials = self.make_credentials() assert not credentials._always_use_jwt_access diff --git a/tests/test_external_account.py b/tests/test_external_account.py index 6f6e18b2c..5225dcf34 100644 --- a/tests/test_external_account.py +++ b/tests/test_external_account.py @@ -505,6 +505,11 @@ def test_universe_domain(self): credentials = self.make_credentials() assert credentials.universe_domain == external_account._DEFAULT_UNIVERSE_DOMAIN + def test_with_universe_domain(self): + credentials = self.make_credentials() + new_credentials = credentials.with_universe_domain("dummy_universe.com") + assert new_credentials.universe_domain == "dummy_universe.com" + def test_info_workforce_pool(self): credentials = self.make_workforce_pool_credentials( workforce_pool_user_project=self.WORKFORCE_POOL_USER_PROJECT