feat: add integration test for universe domain (#1346)

cojenco · web-flow · commit 02a972d35fae · 2024-09-25T03:40:27.000Z
diff --git a/.kokoro/build.sh b/.kokoro/build.sh
@@ -34,6 +34,14 @@ export API_VERSION_OVERRIDE
 export DUAL_REGION_LOC_1
 export DUAL_REGION_LOC_2
 
+# Setup universe domain testing needed environment variables.
+export TEST_UNIVERSE_DOMAIN_CREDENTIAL=$(realpath ${KOKORO_GFILE_DIR}/secret_manager/client-library-test-universe-domain-credential)
+export TEST_UNIVERSE_DOMAIN=$(gcloud secrets versions access latest --project cloud-devrel-kokoro-resources --secret=client-library-test-universe-domain)
+export TEST_UNIVERSE_PROJECT_ID=$(gcloud secrets versions access latest --project cloud-devrel-kokoro-resources --secret=client-library-test-universe-project-id)
+export TEST_UNIVERSE_LOCATION=$(gcloud secrets versions access latest --project cloud-devrel-kokoro-resources --secret=client-library-test-universe-storage-location)
+
+
+
 # Debug: show build environment
 env | grep KOKORO
 
diff --git a/.kokoro/presubmit/system-3.8.cfg b/.kokoro/presubmit/system-3.8.cfg
@@ -4,4 +4,10 @@
 env_vars: {
     key: "NOX_SESSION"
     value: "system-3.8"
+}
+
+# Credentials needed to test universe domain.
+env_vars: {
+  key: "SECRET_MANAGER_KEYS"
+  value: "client-library-test-universe-domain-credential"
 }
diff --git a/owlbot.py b/owlbot.py
@@ -46,6 +46,7 @@
         "noxfile.py",
         "CONTRIBUTING.rst",
         "README.rst",
+        ".kokoro/presubmit/system-3.8.cfg",
         ".kokoro/samples/python3.6", # remove python 3.6 support
         ".github/blunderbuss.yml", # blunderbuss assignment to python squad
         ".github/workflows", # exclude gh actions as credentials are needed for tests
@@ -66,7 +67,15 @@
 
 # Export dual region locations
 export DUAL_REGION_LOC_1
-export DUAL_REGION_LOC_2""")
+export DUAL_REGION_LOC_2
+
+# Setup universe domain testing needed environment variables.
+export TEST_UNIVERSE_DOMAIN_CREDENTIAL=$(realpath ${KOKORO_GFILE_DIR}/secret_manager/client-library-test-universe-domain-credential)
+export TEST_UNIVERSE_DOMAIN=$(gcloud secrets versions access latest --project cloud-devrel-kokoro-resources --secret=client-library-test-universe-domain)
+export TEST_UNIVERSE_PROJECT_ID=$(gcloud secrets versions access latest --project cloud-devrel-kokoro-resources --secret=client-library-test-universe-project-id)
+export TEST_UNIVERSE_LOCATION=$(gcloud secrets versions access latest --project cloud-devrel-kokoro-resources --secret=client-library-test-universe-storage-location)
+
+""")
 
 s.replace(
     ".coveragerc",
diff --git a/tests/system/_helpers.py b/tests/system/_helpers.py
@@ -31,6 +31,10 @@
 
 user_project = os.environ.get("GOOGLE_CLOUD_TESTS_USER_PROJECT")
 testing_mtls = os.getenv("GOOGLE_API_USE_CLIENT_CERTIFICATE") == "true"
+test_universe_domain = os.getenv("TEST_UNIVERSE_DOMAIN")
+test_universe_project_id = os.getenv("TEST_UNIVERSE_PROJECT_ID")
+test_universe_location = os.getenv("TEST_UNIVERSE_LOCATION")
+test_universe_domain_credential = os.getenv("TEST_UNIVERSE_DOMAIN_CREDENTIAL")
 signing_blob_content = b"This time for sure, Rocky!"
 is_api_endpoint_override = (
     _get_default_storage_base_url() != "https://storage.googleapis.com"
diff --git a/tests/system/conftest.py b/tests/system/conftest.py
@@ -331,3 +331,56 @@ def keyring(storage_client, kms_bucket, kms_client):
         except exceptions.NotFound:
             key = {"purpose": purpose}
             kms_client.create_crypto_key(keyring_path, key_name, key)
+
+
+@pytest.fixture(scope="function")
+def test_universe_domain():
+    if _helpers.test_universe_domain is None:
+        pytest.skip("TEST_UNIVERSE_DOMAIN not set in environment.")
+    return _helpers.test_universe_domain
+
+
+@pytest.fixture(scope="function")
+def test_universe_project_id():
+    if _helpers.test_universe_project_id is None:
+        pytest.skip("TEST_UNIVERSE_PROJECT_ID not set in environment.")
+    return _helpers.test_universe_project_id
+
+
+@pytest.fixture(scope="function")
+def test_universe_location():
+    if _helpers.test_universe_location is None:
+        pytest.skip("TEST_UNIVERSE_LOCATION not set in environment.")
+    return _helpers.test_universe_location
+
+
+@pytest.fixture(scope="function")
+def test_universe_domain_credential():
+    if _helpers.test_universe_domain_credential is None:
+        pytest.skip("TEST_UNIVERSE_DOMAIN_CREDENTIAL not set in environment.")
+    return _helpers.test_universe_domain_credential
+
+
+@pytest.fixture(scope="function")
+def universe_domain_credential(test_universe_domain_credential):
+    from google.oauth2 import service_account
+
+    return service_account.Credentials.from_service_account_file(
+        test_universe_domain_credential
+    )
+
+
+@pytest.fixture(scope="function")
+def universe_domain_client(
+    test_universe_domain, test_universe_project_id, universe_domain_credential
+):
+    from google.cloud.storage import Client
+
+    client_options = {"universe_domain": test_universe_domain}
+    ud_storage_client = Client(
+        project=test_universe_project_id,
+        credentials=universe_domain_credential,
+        client_options=client_options,
+    )
+    with contextlib.closing(ud_storage_client):
+        yield ud_storage_client
diff --git a/tests/system/test_client.py b/tests/system/test_client.py
@@ -184,3 +184,30 @@ def test_download_blob_to_file_w_etag(
         if_etag_match=blob.etag,
     )
     assert buffer.getvalue() == payload
+
+
+def test_client_universe_domain(
+    universe_domain_client,
+    test_universe_location,
+    buckets_to_delete,
+    blobs_to_delete,
+):
+    bucket_name = _helpers.unique_name("gcp-systest-ud")
+    ud_bucket = universe_domain_client.create_bucket(
+        bucket_name, location=test_universe_location
+    )
+    buckets_to_delete.append(ud_bucket)
+
+    blob_name = _helpers.unique_name("gcp-systest-ud")
+    blob = ud_bucket.blob(blob_name)
+    payload = b"The quick brown fox jumps over the lazy dog"
+    blob.upload_from_string(payload)
+    blobs_to_delete.append(blob)
+
+    with tempfile.NamedTemporaryFile() as temp_f:
+        with open(temp_f.name, "wb") as file_obj:
+            universe_domain_client.download_blob_to_file(blob, file_obj)
+        with open(temp_f.name, "rb") as file_obj:
+            stored_contents = file_obj.read()
+
+    assert stored_contents == payload
