feat: allow users to configure universe environment variable (#2369)

Author: ohmayr

googleapiclient/discovery.py

@@ -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'.
@@ -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)
 

tests/test_discovery.py

@@ -2600,6 +2600,71 @@ 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()