Skip to content

Commit

Permalink
feat: allow users to configure universe environment variable
Browse files Browse the repository at this point in the history
  • Loading branch information
ohmayr committed Mar 25, 2024
1 parent 359b3ac commit e08e440
Show file tree
Hide file tree
Showing 2 changed files with 64 additions and 1 deletion.
4 changes: 3 additions & 1 deletion googleapiclient/discovery.py
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,7 @@
# Parameters controlling mTLS behavior. See https://google.aip.dev/auth/4114.
GOOGLE_API_USE_CLIENT_CERTIFICATE = "GOOGLE_API_USE_CLIENT_CERTIFICATE"
GOOGLE_API_USE_MTLS_ENDPOINT = "GOOGLE_API_USE_MTLS_ENDPOINT"
GOOGLE_CLOUD_UNIVERSE_DOMAIN = "GOOGLE_CLOUD_UNIVERSE_DOMAIN"

# Parameters accepted by the stack, but not visible via discovery.
# TODO(dhermes): Remove 'userip' in 'v2'.
Expand Down Expand Up @@ -553,8 +554,9 @@ def build_from_document(
base = urllib.parse.urljoin(service["rootUrl"], service["servicePath"])
universe_domain = None
if HAS_UNIVERSE:
universe_domain_env = os.getenv(GOOGLE_CLOUD_UNIVERSE_DOMAIN, None)
universe_domain = universe.determine_domain(
client_options.universe_domain, None
client_options.universe_domain, universe_domain_env
)
base = base.replace(universe.DEFAULT_UNIVERSE, universe_domain)

Expand Down
61 changes: 61 additions & 0 deletions tests/test_discovery.py
Original file line number Diff line number Diff line change
Expand Up @@ -2600,6 +2600,67 @@ def test_client_options_universe_configured_with_api_override(self):

assert tasks._baseUrl == fake_api_endpoint

def test_universe_env_var_configured_empty(self):
credentials = mock.Mock(spec=google.auth.credentials.Credentials)
discovery = read_datafile("tasks.json")

with self.assertRaises(universe.EmptyUniverseError):
with mock.patch.dict(
"os.environ", {"GOOGLE_CLOUD_UNIVERSE_DOMAIN": ""}
):
tasks = build_from_document(
discovery,
credentials=credentials,
)

def test_universe_env_var_configured_with_mtls(self):
fake_universe = "foo.com"
discovery = read_datafile("tasks.json")

with self.assertRaises(MutualTLSChannelError):
with mock.patch.dict(
"os.environ", {"GOOGLE_API_USE_MTLS_ENDPOINT": "always",
"GOOGLE_CLOUD_UNIVERSE_DOMAIN": fake_universe}
):
tasks = build_from_document(discovery)

def test_universe_env_var_configured_with_api_override(self):
fake_universe = "foo.com"
fake_api_endpoint = "https://www.bar.com/"
credentials = mock.Mock(universe_domain=fake_universe)
discovery = read_datafile("tasks.json")

with mock.patch.dict(
"os.environ", {"GOOGLE_CLOUD_UNIVERSE_DOMAIN": fake_universe}
):
tasks = build_from_document(
discovery,
credentials=credentials,
client_options=google.api_core.client_options.ClientOptions(
api_endpoint=fake_api_endpoint
),
)

assert tasks._baseUrl == fake_api_endpoint

def test_universe_env_var_configured_with_client_options_universe(self):
fake_universe = "foo.com"
another_fake_universe = "bar.com"
credentials = mock.Mock(universe_domain=fake_universe)
discovery = read_datafile("tasks.json")

with mock.patch.dict(
"os.environ", {"GOOGLE_CLOUD_UNIVERSE_DOMAIN": another_fake_universe}
):
tasks = build_from_document(
discovery,
credentials=credentials,
client_options=google.api_core.client_options.ClientOptions(
universe_domain=fake_universe
),
)

assert tasks._universe_domain == fake_universe

if __name__ == "__main__":
unittest.main()

0 comments on commit e08e440

Please sign in to comment.