Skip to content

Commit bcd3cb6

Browse files
twishabansaltwishabansal
committed
added e2e tests
1 parent f53a4da commit bcd3cb6

File tree

3 files changed

+231
-1
lines changed

3 files changed

+231
-1
lines changed

‎packages/toolbox-core/pyproject.toml

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,9 @@ test = [
4545
"mypy==1.15.0",
4646
"pytest==8.3.5",
4747
"pytest-aioresponses==0.3.0",
48-
"pytest-asyncio==0.26.0"
48+
"pytest-asyncio==0.26.0",
49+
"google-cloud-secret-manager==2.23.2",
50+
"google-cloud-storage==3.1.0",
4951
]
5052
[build-system]
5153
requires = ["setuptools"]

‎packages/toolbox-core/tests/conftest.py

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
# Copyright 2024 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()

‎packages/toolbox-core/tests/test_e2e.py

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
# Copyright 2024 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+
import pytest
16+
import pytest_asyncio
17+
18+
from toolbox_core.client import ToolboxClient
19+
from toolbox_core.tool import ToolboxTool
20+
21+
22+
@pytest.mark.asyncio
23+
@pytest.mark.usefixtures("toolbox_server")
24+
class TestE2EClient:
25+
@pytest_asyncio.fixture(scope="function")
26+
async def toolbox(self):
27+
toolbox = ToolboxClient("http://localhost:5000")
28+
return toolbox
29+
30+
@pytest_asyncio.fixture(scope="function")
31+
async def get_n_rows_tool(self, toolbox: ToolboxClient) -> ToolboxTool:
32+
tool = await toolbox.load_tool("get-n-rows")
33+
assert tool.__name__ == "get-n-rows"
34+
return tool
35+
36+
#### Basic e2e tests
37+
@pytest.mark.parametrize(
38+
"toolset_name, expected_length, expected_tools",
39+
[
40+
("my-toolset", 1, ["get-row-by-id"]),
41+
("my-toolset-2", 2, ["get-n-rows", "get-row-by-id"]),
42+
],
43+
)
44+
async def test_load_toolset_specific(
45+
self,
46+
toolbox: ToolboxClient,
47+
toolset_name: str,
48+
expected_length: int,
49+
expected_tools: list[str],
50+
):
51+
toolset = await toolbox.load_toolset(toolset_name)
52+
assert len(toolset) == expected_length
53+
tool_names = {tool.__name__ for tool in toolset}
54+
assert tool_names == set(expected_tools)
55+
56+
async def test_run_tool(self, get_n_rows_tool: ToolboxTool):
57+
response = await get_n_rows_tool(num_rows="2")
58+
59+
assert isinstance(response, str)
60+
assert "row1" in response
61+
assert "row2" in response
62+
assert "row3" not in response

0 commit comments

Comments
 (0)