Skip to content

Commit 1c60ee2

Browse files
sararobcopybara-github
sararob
authored and
copybara-github
committed
chore: Enable replay tests for Vertex SDK GenAI client
PiperOrigin-RevId: 774817886
1 parent f63e436 commit 1c60ee2

File tree

5 files changed

+411
-5
lines changed

5 files changed

+411
-5
lines changed

‎tests/unit/vertexai/genai/replays/conftest.py

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
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+
16+
17+
"""Conftest for Vertex SDK GenAI tests."""
18+
19+
import os
20+
from unittest import mock
21+
22+
from vertexai._genai import (
23+
client as vertexai_genai_client_module,
24+
)
25+
from google.genai import _replay_api_client
26+
from google.genai import client as google_genai_client_module
27+
import pytest
28+
29+
30+
def pytest_addoption(parser):
31+
parser.addoption(
32+
"--mode",
33+
action="store",
34+
default="auto",
35+
help="""Replay mode.
36+
One of:
37+
* auto: Replay if replay files exist, otherwise record.
38+
* record: Always call the API and record.
39+
* replay: Always replay, fail if replay files do not exist.
40+
* api: Always call the API and do not record.
41+
* tap: Always replay, fail if replay files do not exist. Also sets default values for the API key and replay directory.
42+
""",
43+
)
44+
parser.addoption(
45+
"--replays-directory-prefix",
46+
action="store",
47+
default=None,
48+
help="""Directory to use for replays.
49+
If not set, the default directory will be used.
50+
""",
51+
)
52+
53+
54+
@pytest.fixture
55+
def use_vertex():
56+
return True
57+
58+
59+
# Overridden at the module level for each test file.
60+
@pytest.fixture
61+
def replays_prefix():
62+
return "test"
63+
64+
65+
def _get_replay_id(use_vertex: bool, replays_prefix: str) -> str:
66+
test_name_ending = os.environ.get("PYTEST_CURRENT_TEST").split("::")[-1]
67+
test_name = test_name_ending.split(" ")[0].split("[")[0] + "." + "vertex"
68+
return "/".join([replays_prefix, test_name])
69+
70+
71+
@pytest.fixture
72+
def client(use_vertex, replays_prefix, http_options, request):
73+
74+
mode = request.config.getoption("--mode")
75+
replays_directory_prefix = request.config.getoption("--replays-directory-prefix")
76+
if mode not in ["auto", "record", "replay", "api", "tap"]:
77+
raise ValueError("Invalid mode: " + mode)
78+
test_function_name = request.function.__name__
79+
test_filename = os.path.splitext(os.path.basename(request.path))[0]
80+
if test_function_name.startswith(test_filename):
81+
raise ValueError(
82+
f"""
83+
{test_function_name}:
84+
Do not include the test filename in the test function name.
85+
keep the test function name short."""
86+
)
87+
88+
replay_id = _get_replay_id(use_vertex, replays_prefix)
89+
90+
if mode == "tap":
91+
mode = "replay"
92+
93+
# Set various environment variables to ensure that the test runs.
94+
os.environ["GOOGLE_API_KEY"] = "dummy-api-key"
95+
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.join(
96+
os.path.dirname(__file__),
97+
"credentials.json",
98+
)
99+
os.environ["GOOGLE_CLOUD_PROJECT"] = "project-id"
100+
os.environ["GOOGLE_CLOUD_LOCATION"] = "location"
101+
102+
# Set the replay directory to the root directory of the replays.
103+
# This is needed to ensure that the replay files are found.
104+
replays_root_directory = os.path.abspath(
105+
os.path.join(
106+
os.path.dirname(__file__),
107+
"../../../../../../../../../google/cloud/aiplatform/sdk/genai/replays",
108+
)
109+
)
110+
os.environ["GOOGLE_GENAI_REPLAYS_DIRECTORY"] = replays_root_directory
111+
replay_client = _replay_api_client.ReplayApiClient(
112+
mode=mode,
113+
replay_id=replay_id,
114+
vertexai=use_vertex,
115+
http_options=http_options,
116+
)
117+
118+
replay_client.replays_directory = (
119+
f"{replays_directory_prefix}/google/cloud/aiplatform/sdk/replays/"
120+
)
121+
122+
with mock.patch.object(
123+
google_genai_client_module.Client, "_get_api_client"
124+
) as patch_method:
125+
patch_method.return_value = replay_client
126+
google_genai_client = vertexai_genai_client_module.Client()
127+
128+
# Yield the client so that cleanup can be completed at the end of the test.
129+
yield google_genai_client
130+
131+
# Save the replay after the test if we're in recording mode.
132+
replay_client.close()

‎tests/unit/vertexai/genai/replays/pytest_helper.py

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
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+
16+
import os
17+
from typing import Any, Optional
18+
19+
from google.genai._api_client import HttpOptions
20+
import pytest
21+
22+
23+
is_api_mode = "config.getoption('--mode') == 'api'"
24+
25+
26+
# Sets up the test framework.
27+
# file: Always use __file__
28+
# globals_for_file: Always use globals()
29+
def setup(
30+
*,
31+
file: str,
32+
globals_for_file: Optional[dict[str, Any]] = None,
33+
test_method: Optional[str] = None,
34+
http_options: Optional[HttpOptions] = None,
35+
):
36+
"""Generates parameterization for tests"""
37+
replays_directory = (
38+
file.replace(os.path.dirname(__file__), "tests/vertex_sdk_genai_replays")
39+
.replace(".py", "")
40+
.replace("/test_", "/")
41+
)
42+
43+
# Add fixture for requested client option.
44+
return pytest.mark.parametrize(
45+
"use_vertex, replays_prefix, http_options",
46+
[
47+
(True, replays_directory, http_options),
48+
],
49+
)

‎tests/unit/vertexai/genai/replays/test_evaluate_instances.py

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
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+
# pylint: disable=protected-access,bad-continuation,missing-function-docstring
16+
17+
import os
18+
19+
from tests.unit.vertexai.genai.replays import pytest_helper
20+
from vertexai._genai import types
21+
import pytest
22+
23+
24+
IS_KOKORO = os.getenv("KOKORO_BUILD_NUMBER") is not None
25+
26+
27+
@pytest.mark.skipif(IS_KOKORO, reason="This test is only run in google3 env.")
28+
class TestEvaluateInstances:
29+
"""Tests for evaluate instances."""
30+
31+
def test_bleu_metric(self, client):
32+
test_bleu_input = types.BleuInput(
33+
instances=[
34+
types.BleuInstance(
35+
reference="The quick brown fox jumps over the lazy dog.",
36+
prediction="A fast brown fox leaps over a lazy dog.",
37+
)
38+
],
39+
metric_spec=types.BleuSpec(),
40+
)
41+
response = client.evals._evaluate_instances(bleu_input=test_bleu_input)
42+
assert len(response.bleu_results.bleu_metric_values) == 1
43+
44+
45+
pytestmark = pytest_helper.setup(
46+
file=__file__,
47+
globals_for_file=globals(),
48+
test_method="evals.evaluate",
49+
)

‎tests/unit/vertexai/genai/run_replay_tests.sh

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
#!/bin/bash
2+
3+
# This script runs replay tests for the Vertex SDK GenAI client
4+
# It is intended to be used from the google3 directory of a CitC client.
5+
# You can provide a specific test file to run, or it will run all the replay tests
6+
# in third_party/py/google/cloud/aiplatform/tests/unit/vertexai/genai/replays/
7+
#
8+
# Example:
9+
# ./third_party/py/google/cloud/aiplatform/tests/unit/vertexai/genai/run_replay_tests.sh test_evals.py
10+
11+
# It also supports a --mode flag, which can be one of:
12+
# * record: Call the API and record the result in a replay file.
13+
# * replay: Use the recorded replay file to simulate the API call, or record if the replay file does not exist.
14+
# * api: Call the API and do not record.
15+
16+
# Get the current working directory
17+
START_DIR=$(pwd)
18+
19+
# Check if the current directory ends with /google3
20+
# Otherwise the copybara command will fail
21+
if [[ "$START_DIR" != */google3 ]]; then
22+
echo "Error: This script must be run from your client's '/google3' directory."
23+
echo "Your current directory is: $START_DIR"
24+
exit 1
25+
fi
26+
27+
# Check required env vars have been set
28+
REQUIRED_ENV_VARS=(
29+
"GOOGLE_CLOUD_PROJECT"
30+
"GOOGLE_CLOUD_LOCATION"
31+
"GOOGLE_GENAI_REPLAYS_DIRECTORY"
32+
)
33+
34+
for var_name in "${REQUIRED_ENV_VARS[@]}"; do
35+
var_value="${!var_name}"
36+
if [ -z "$var_value" ]; then
37+
echo "Error: Environment variable $var_name is not set."
38+
echo "Please set it before running this script."
39+
exit 1
40+
fi
41+
done
42+
43+
# Generate a unique temporary directory in /tmp/
44+
TEMP_DIR=$(mktemp -d -t XXXXXX)
45+
46+
# Check if the temporary directory was created successfully
47+
if [ -z "$TEMP_DIR" ]; then
48+
echo "Error: Could not create a temporary directory."
49+
exit 1
50+
fi
51+
52+
echo "Created temporary directory: $TEMP_DIR"
53+
54+
# Run copybara and copy Vertex SDK to the temporary directory
55+
# The --folder-dir argument is set to the newly created temporary directory.
56+
echo "Running copybara..."
57+
COPYBARA_EXEC="/google/bin/releases/copybara/public/copybara/copybara"
58+
"$COPYBARA_EXEC" third_party/py/google/cloud/aiplatform/copy.bara.sky folder_to_folder .. --folder-dir="$TEMP_DIR" --ignore-noop
59+
60+
# Check copybara's exit status
61+
if [ $? -ne 0 ]; then
62+
echo "Error: copybara command failed. Exiting."
63+
# Clean up the temporary directory on failure
64+
rm -rf "$TEMP_DIR"
65+
exit 1
66+
fi
67+
68+
echo "Copybara completed successfully."
69+
70+
# Change into the temporary directory with copybara output
71+
echo "Changing into temp directory: $TEMP_DIR"
72+
cd "$TEMP_DIR"
73+
74+
if [ $? -ne 0 ]; then
75+
echo "Error: Could not change into directory $TEMP_DIR. Exiting."
76+
exit 1
77+
fi
78+
79+
PARSED_ARGS=$(getopt -o "" -l "mode:" -- "$@")
80+
81+
if [ $? -ne 0 ]; then
82+
echo "Error: Failed to parse command line arguments." >&2
83+
exit 1
84+
fi
85+
86+
# Get the test file path and --mode flag value if provided
87+
eval set -- "$PARSED_ARGS"
88+
89+
TEST_FILE_ARG="" # Stores the provided test path, if any
90+
MODE_VALUE="" # Stores the value of the --mode flag (e.g., 'replay')
91+
92+
while true; do
93+
case "$1" in
94+
--mode)
95+
MODE_VALUE="$2"
96+
shift 2
97+
;;
98+
--)
99+
shift
100+
break
101+
;;
102+
*)
103+
echo "Internal error: Unrecognized arg option: '$1'" >&2
104+
exit 1
105+
;;
106+
esac
107+
done
108+
109+
# We expect at most one positional argument (the test file path).
110+
if [ -n "$1" ]; then
111+
TEST_FILE_ARG="$1"
112+
if [ "$#" -gt 1 ]; then
113+
echo "Warning: Ignoring extra positional arguments after '$TEST_FILE_ARG'. Only one test file/path can be specified." >&2
114+
fi
115+
fi
116+
117+
# Construct the full --mode argument string to be passed to pytest.
118+
MODE_ARG=""
119+
if [ -n "$MODE_VALUE" ]; then
120+
MODE_ARG="--mode $MODE_VALUE"
121+
fi
122+
123+
124+
# Set pytest path for which tests to run
125+
DEFAULT_TEST_PATH="tests/unit/vertexai/genai/replays/"
126+
127+
if [ -n "$TEST_FILE_ARG" ]; then
128+
PYTEST_PATH="${DEFAULT_TEST_PATH}${TEST_FILE_ARG}"
129+
echo "Provided test file path: '$TEST_FILE_ARG'. Running pytest on: ${PYTEST_PATH}"
130+
else
131+
PYTEST_PATH="$DEFAULT_TEST_PATH"
132+
echo "No test file arg provided. Running pytest on default path: ${PYTEST_PATH}"
133+
fi
134+
135+
# Run tests
136+
# -s is equivalent to --capture=no, it ensures pytest doesn't capture the output from stdout and stderr
137+
# so it can be logged when this script is run
138+
pytest -v -s "$PYTEST_PATH" ${MODE_ARG} --replays-directory-prefix="$START_DIR"
139+
140+
PYTEST_EXIT_CODE=$?
141+
142+
echo "Cleaning up temporary directory: $TEMP_DIR"
143+
# Go back to the original directory before removing the temporary directory
144+
cd - > /dev/null
145+
rm -rf "$TEMP_DIR"
146+
147+
echo "Pytest tests completed with exit code: $PYTEST_EXIT_CODE."
148+
149+
exit $PYTEST_EXIT_CODE

0 commit comments

Comments
 (0)