247 lines
8.1 KiB
Python
247 lines
8.1 KiB
Python
# Copyright 2025 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
|
|
import contextlib
|
|
import json
|
|
import os
|
|
import pathlib
|
|
from typing import Any, Optional
|
|
from pydantic import BaseModel, Field, SerializeAsAny
|
|
import pytest
|
|
import re
|
|
from .. import _common
|
|
from .. import _replay_api_client
|
|
from .. import types
|
|
from .._api_client import HttpOptions
|
|
|
|
is_api_mode = "config.getoption('--mode') == 'api'"
|
|
|
|
|
|
class TestTableItem(types.TestTableItem):
|
|
# This is not a test suite class.
|
|
__test__ = False
|
|
|
|
# Overrides to support Pydantic models.
|
|
parameters: SerializeAsAny[BaseModel] = Field(
|
|
description="""The parameters to the test. Use pydantic models.""",
|
|
)
|
|
|
|
|
|
def base_test_function(
|
|
client,
|
|
use_vertex: bool,
|
|
replays_prefix: str,
|
|
test_method: str,
|
|
test_table_item: TestTableItem,
|
|
globals_for_file: dict[str, Any],
|
|
):
|
|
replay_id = (
|
|
test_table_item.override_replay_id
|
|
if test_table_item.override_replay_id
|
|
else test_table_item.name
|
|
)
|
|
api_type = 'vertex' if use_vertex else 'mldev'
|
|
replay_id = f'{replays_prefix}/{replay_id}.{api_type}'
|
|
client._api_client.initialize_replay_session(replay_id)
|
|
# vars().copy() provides a shallow copy of the parameters.
|
|
parameters_dict = vars(test_table_item.parameters).copy()
|
|
try:
|
|
if '.' in test_method:
|
|
method_name_parts = test_method.split('.')
|
|
module_path_parts = method_name_parts[:-1]
|
|
method_name = method_name_parts[-1]
|
|
# Iterates the nested module starting from client
|
|
current_object = client
|
|
for part in module_path_parts:
|
|
current_object = getattr(current_object, part)
|
|
method = getattr(current_object, method_name)
|
|
method(**parameters_dict)
|
|
else:
|
|
custom_method = globals_for_file[test_method]
|
|
custom_method(client, test_table_item.parameters)
|
|
# Should not reach here if expecting an exception.
|
|
if test_table_item.exception_if_mldev and not client._api_client.vertexai:
|
|
assert False, 'Should have raised exception in MLDev.'
|
|
elif test_table_item.exception_if_vertex and client._api_client.vertexai:
|
|
assert False, 'Should have raised exception in Vertex.'
|
|
client._api_client.close()
|
|
except Exception as e:
|
|
if test_table_item.exception_if_mldev and not client._api_client.vertexai:
|
|
if test_table_item.exception_if_mldev not in str(e):
|
|
raise AssertionError(
|
|
f"'{test_table_item.exception_if_mldev}' not in '{str(e)}'"
|
|
) from e
|
|
elif test_table_item.exception_if_vertex and client._api_client.vertexai:
|
|
if test_table_item.exception_if_vertex not in str(e):
|
|
raise AssertionError(
|
|
f"'{test_table_item.exception_if_vertex}' not in '{str(e)}'"
|
|
) from e
|
|
else:
|
|
raise e
|
|
|
|
|
|
def create_test_for_table_item(
|
|
globals_for_file: dict[str, Any],
|
|
test_method: str,
|
|
test_table_item: TestTableItem,
|
|
):
|
|
return lambda client, use_vertex, replays_prefix: base_test_function(
|
|
client,
|
|
use_vertex,
|
|
replays_prefix,
|
|
test_method,
|
|
test_table_item,
|
|
globals_for_file,
|
|
)
|
|
|
|
|
|
def create_test_for_table(
|
|
globals_for_file: dict[str, Any],
|
|
test_method: str,
|
|
test_table: list[TestTableItem],
|
|
):
|
|
for test_table_item in test_table:
|
|
if test_table_item.has_union:
|
|
assert test_table_item.name.startswith('test_union_'), f"""
|
|
test: {test_table_item.name}
|
|
When has_union is true, it must be a test that explicitly test the
|
|
non-canonical type in the union.
|
|
For all other tests, use transformers to convert data to the canonical
|
|
type of the field so it can be tested in other languages.
|
|
If this test is truly about testing the non-canonical type of the union,
|
|
rename the test to start with 'test_union_'. E.g.
|
|
test_union_contents_is_string()"""
|
|
globals_for_file[test_table_item.name] = create_test_for_table_item(
|
|
globals_for_file, test_method, test_table_item
|
|
)
|
|
|
|
|
|
# Sets up the test framework.
|
|
# file: Always use __file__
|
|
# globals_for_file: Always use globals()
|
|
# test_table: The test table for the file.
|
|
# Use test table over individual tests whenever possible.
|
|
# Tests built with test_table will run in all other languages automatically.
|
|
# Otherwise, you will need to write tests in all other languages manually.
|
|
def setup(
|
|
*,
|
|
file: str,
|
|
globals_for_file: Optional[dict[str, Any]] = None,
|
|
test_method: Optional[str] = None,
|
|
test_table: Optional[list[TestTableItem]] = None,
|
|
http_options: Optional[HttpOptions] = None,
|
|
):
|
|
"""Generates parameterization for tests, run for both Vertex and MLDev."""
|
|
replays_directory = (
|
|
file.replace(os.path.dirname(__file__), 'tests')
|
|
.replace('.py', '')
|
|
.replace('/test_', '/')
|
|
)
|
|
|
|
is_tap_mode = False
|
|
if os.environ.get('UNITTEST_ON_FORGE', None) is not None:
|
|
is_tap_mode = True
|
|
|
|
if test_table:
|
|
create_test_for_table(globals_for_file, test_method, test_table)
|
|
|
|
if test_table and not is_tap_mode:
|
|
replays_root_directory = os.environ.get(
|
|
'GOOGLE_GENAI_REPLAYS_DIRECTORY', None
|
|
)
|
|
if replays_root_directory is None:
|
|
raise ValueError(
|
|
'GOOGLE_GENAI_REPLAYS_DIRECTORY environment variable is not set'
|
|
)
|
|
abs_replay_directory = os.path.join(
|
|
replays_root_directory, replays_directory
|
|
)
|
|
test_table_file_path = os.path.join(
|
|
abs_replay_directory, '_test_table.json'
|
|
)
|
|
|
|
pathlib.Path(abs_replay_directory).mkdir(parents=True, exist_ok=True)
|
|
assert isinstance(
|
|
test_table[0].parameters, BaseModel
|
|
), f'{test_table_file_path} parameters must be a BaseModel.'
|
|
test_table_file = types.TestTableFile(
|
|
comment='Auto-generated. Do not edit.',
|
|
test_method=test_method,
|
|
parameter_names=list(test_table[0].parameters.model_fields.keys()),
|
|
test_table=test_table,
|
|
)
|
|
os.makedirs(os.path.dirname(test_table_file_path), exist_ok=True)
|
|
|
|
with open(test_table_file_path, 'w') as f:
|
|
f.write(
|
|
test_table_file.model_dump_json(exclude_none=True, by_alias=True, indent=2),
|
|
)
|
|
|
|
# Add fixture for requested client option.
|
|
return pytest.mark.parametrize(
|
|
'use_vertex, replays_prefix, http_options',
|
|
[
|
|
(True, replays_directory, http_options),
|
|
(False, replays_directory, http_options),
|
|
],
|
|
)
|
|
|
|
|
|
def exception_if_mldev(client, exception_type: type[Exception]):
|
|
if client._api_client.vertexai:
|
|
return contextlib.nullcontext()
|
|
else:
|
|
return pytest.raises(exception_type)
|
|
|
|
|
|
def exception_if_vertex(client, exception_type: type[Exception]):
|
|
if client._api_client.vertexai:
|
|
return pytest.raises(exception_type)
|
|
else:
|
|
return contextlib.nullcontext()
|
|
|
|
|
|
def snake_to_camel(snake_str: str) -> str:
|
|
"""Converts a snake_case string to CamelCase."""
|
|
return re.sub(r'_([a-zA-Z])', lambda match: match.group(1).upper(), snake_str)
|
|
|
|
|
|
def camel_to_snake(camel_str: str) -> str:
|
|
"""Converts a CamelCase string to snake_case."""
|
|
return re.sub(r'([A-Z])', r'_\1', camel_str).lower().lstrip('_')
|
|
|
|
|
|
def get_value_ignore_key_case(obj, key):
|
|
"""Returns the value of the key in the object, converting to camelCase or snake_case if necessary."""
|
|
return obj.get(snake_to_camel(key), obj.get(camel_to_snake(key), None))
|
|
|
|
|
|
def camel_to_snake_all_keys(data):
|
|
"""Converts all keys in a dictionary or list to snake_case."""
|
|
if isinstance(data, dict):
|
|
new_dict = {}
|
|
for key, value in data.items():
|
|
if isinstance(key, str):
|
|
new_key = camel_to_snake(key)
|
|
else:
|
|
new_key = key
|
|
new_dict[new_key] = camel_to_snake_all_keys(value)
|
|
return new_dict
|
|
elif isinstance(data, list):
|
|
return [camel_to_snake_all_keys(item) for item in data]
|
|
else:
|
|
return data
|