|
| 1 | +# Copyright 2025 Google LLC |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +"""Contains pytest fixtures that are accessible from all |
| 16 | +files present in the same directory.""" |
| 17 | + |
| 18 | +from __future__ import annotations |
| 19 | + |
| 20 | +import os |
| 21 | +import platform |
| 22 | +import subprocess |
| 23 | +import tempfile |
| 24 | +import time |
| 25 | +from typing import Generator |
| 26 | + |
| 27 | +import google |
| 28 | +import pytest_asyncio |
| 29 | +from google.auth import compute_engine |
| 30 | +from google.cloud import secretmanager, storage |
| 31 | + |
| 32 | + |
| 33 | +#### Define Utility Functions |
| 34 | +def get_env_var(key: str) -> str: |
| 35 | + """Gets environment variables.""" |
| 36 | + value = os.environ.get(key) |
| 37 | + if value is None: |
| 38 | + raise ValueError(f"Must set env var {key}") |
| 39 | + return value |
| 40 | + |
| 41 | + |
| 42 | +def access_secret_version( |
| 43 | + project_id: str, secret_id: str, version_id: str = "latest" |
| 44 | +) -> str: |
| 45 | + """Accesses the payload of a given secret version from Secret Manager.""" |
| 46 | + client = secretmanager.SecretManagerServiceClient() |
| 47 | + name = f"projects/{project_id}/secrets/{secret_id}/versions/{version_id}" |
| 48 | + response = client.access_secret_version(request={"name": name}) |
| 49 | + return response.payload.data.decode("UTF-8") |
| 50 | + |
| 51 | + |
| 52 | +def create_tmpfile(content: str) -> str: |
| 53 | + """Creates a temporary file with the given content.""" |
| 54 | + with tempfile.NamedTemporaryFile(delete=False, mode="w") as tmpfile: |
| 55 | + tmpfile.write(content) |
| 56 | + return tmpfile.name |
| 57 | + |
| 58 | + |
| 59 | +def download_blob( |
| 60 | + bucket_name: str, source_blob_name: str, destination_file_name: str |
| 61 | +) -> None: |
| 62 | + """Downloads a blob from a GCS bucket.""" |
| 63 | + storage_client = storage.Client() |
| 64 | + |
| 65 | + bucket = storage_client.bucket(bucket_name) |
| 66 | + blob = bucket.blob(source_blob_name) |
| 67 | + blob.download_to_filename(destination_file_name) |
| 68 | + |
| 69 | + print(f"Blob {source_blob_name} downloaded to {destination_file_name}.") |
| 70 | + |
| 71 | + |
| 72 | +def get_toolbox_binary_url(toolbox_version: str) -> str: |
| 73 | + """Constructs the GCS path to the toolbox binary.""" |
| 74 | + os_system = platform.system().lower() |
| 75 | + arch = ( |
| 76 | + "arm64" if os_system == "darwin" and platform.machine() == "arm64" else "amd64" |
| 77 | + ) |
| 78 | + return f"v{toolbox_version}/{os_system}/{arch}/toolbox" |
| 79 | + |
| 80 | + |
| 81 | +def get_auth_token(client_id: str) -> str: |
| 82 | + """Retrieves an authentication token""" |
| 83 | + request = google.auth.transport.requests.Request() |
| 84 | + credentials = compute_engine.IDTokenCredentials( |
| 85 | + request=request, |
| 86 | + target_audience=client_id, |
| 87 | + use_metadata_identity_endpoint=True, |
| 88 | + ) |
| 89 | + if not credentials.valid: |
| 90 | + credentials.refresh(request) |
| 91 | + return credentials.token |
| 92 | + |
| 93 | + |
| 94 | +#### Define Fixtures |
| 95 | +@pytest_asyncio.fixture(scope="session") |
| 96 | +def project_id() -> str: |
| 97 | + return get_env_var("GOOGLE_CLOUD_PROJECT") |
| 98 | + |
| 99 | + |
| 100 | +@pytest_asyncio.fixture(scope="session") |
| 101 | +def toolbox_version() -> str: |
| 102 | + return get_env_var("TOOLBOX_VERSION") |
| 103 | + |
| 104 | + |
| 105 | +@pytest_asyncio.fixture(scope="session") |
| 106 | +def tools_file_path(project_id: str) -> Generator[str]: |
| 107 | + """Provides a temporary file path containing the tools manifest.""" |
| 108 | + tools_manifest = access_secret_version( |
| 109 | + project_id=project_id, secret_id="sdk_testing_tools" |
| 110 | + ) |
| 111 | + tools_file_path = create_tmpfile(tools_manifest) |
| 112 | + yield tools_file_path |
| 113 | + os.remove(tools_file_path) |
| 114 | + |
| 115 | + |
| 116 | +@pytest_asyncio.fixture(scope="session") |
| 117 | +def auth_token1(project_id: str) -> str: |
| 118 | + client_id = access_secret_version( |
| 119 | + project_id=project_id, secret_id="sdk_testing_client1" |
| 120 | + ) |
| 121 | + return get_auth_token(client_id) |
| 122 | + |
| 123 | + |
| 124 | +@pytest_asyncio.fixture(scope="session") |
| 125 | +def auth_token2(project_id: str) -> str: |
| 126 | + client_id = access_secret_version( |
| 127 | + project_id=project_id, secret_id="sdk_testing_client2" |
| 128 | + ) |
| 129 | + return get_auth_token(client_id) |
| 130 | + |
| 131 | + |
| 132 | +@pytest_asyncio.fixture(scope="session") |
| 133 | +def toolbox_server(toolbox_version: str, tools_file_path: str) -> Generator[None]: |
| 134 | + """Starts the toolbox server as a subprocess.""" |
| 135 | + print("Downloading toolbox binary from gcs bucket...") |
| 136 | + source_blob_name = get_toolbox_binary_url(toolbox_version) |
| 137 | + download_blob("genai-toolbox", source_blob_name, "toolbox") |
| 138 | + print("Toolbox binary downloaded successfully.") |
| 139 | + try: |
| 140 | + print("Opening toolbox server process...") |
| 141 | + # Make toolbox executable |
| 142 | + os.chmod("toolbox", 0o700) |
| 143 | + # Run toolbox binary |
| 144 | + toolbox_server = subprocess.Popen( |
| 145 | + ["./toolbox", "--tools_file", tools_file_path] |
| 146 | + ) |
| 147 | + |
| 148 | + # Wait for server to start |
| 149 | + # Retry logic with a timeout |
| 150 | + for _ in range(5): # retries |
| 151 | + time.sleep(4) |
| 152 | + print("Checking if toolbox is successfully started...") |
| 153 | + if toolbox_server.poll() is None: |
| 154 | + print("Toolbox server started successfully.") |
| 155 | + break |
| 156 | + else: |
| 157 | + raise RuntimeError("Toolbox server failed to start after 5 retries.") |
| 158 | + except subprocess.CalledProcessError as e: |
| 159 | + print(e.stderr.decode("utf-8")) |
| 160 | + print(e.stdout.decode("utf-8")) |
| 161 | + raise RuntimeError(f"{e}\n\n{e.stderr.decode('utf-8')}") from e |
| 162 | + yield |
| 163 | + |
| 164 | + # Clean up toolbox server |
| 165 | + toolbox_server.terminate() |
| 166 | + toolbox_server.wait() |
0 commit comments