feat(phase1): add image storage utilities

- Create backend/services/image_storage.py with 4 core functions:
  - sanitize_filename(): remove unsafe chars, limit to 255 chars, convert to lowercase
  - get_unique_filename(): handle collisions with UUID suffix (format: {name}_{uuid8}_{variant}.jpg)
  - ensure_image_directories(): create /images/ root and category subdirs on startup
  - save_image(): save bytes to /images/{category}/{filename}, returns relative path
- Create comprehensive test suite (22 tests) covering all functionality
- Integrate ensure_image_directories() into FastAPI startup event
- Directory structure: /images/{category}/{filename}
- Collision handling: auto-suffix with UUID if filename exists
- All tests passing, pathlib.Path for safe operations
This commit is contained in:
2026-04-20 21:57:26 +03:00
parent 39fab336ba
commit ea49cd6e4a
4487 changed files with 1305959 additions and 4 deletions

View File

@@ -0,0 +1,21 @@
# 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.
#
"""Tests for the Google GenAI SDK."""
import pytest
pytest.register_assert_rewrite('genai.replay_api_client')

View File

@@ -0,0 +1,21 @@
# 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.
#
"""AFC helper function unit tests.
For AFC end to end test, please write test cases in
test_generate_content_tools.py module
"""

View File

@@ -0,0 +1,309 @@
# 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.
#
"""Test convert_if_exist_pydantic_model."""
import inspect
from typing import Optional, Union
import pydantic
import pytest
import sys
from ... import errors
from ..._extra_utils import convert_if_exist_pydantic_model
def test_builtin_types():
assert convert_if_exist_pydantic_model(1, int, 'param_name', 'func_name') == 1
assert (
convert_if_exist_pydantic_model(1.0, float, 'param_name', 'func_name')
== 1.0
)
assert (
convert_if_exist_pydantic_model('1.0', str, 'param_name', 'func_name')
== '1.0'
)
assert (
convert_if_exist_pydantic_model(True, bool, 'param_name', 'func_name')
is True
)
assert convert_if_exist_pydantic_model(
[1], list, 'param_name', 'func_name'
) == [1]
assert convert_if_exist_pydantic_model(
{'key1': 1}, dict, 'param_name', 'func_name'
) == {'key1': 1}
assert convert_if_exist_pydantic_model(
{'key1': 1, 'key2': 2}, dict[str, int], 'param_name', 'func_name'
) == {'key1': 1, 'key2': 2}
def test_value_int_annotation_float():
assert (
convert_if_exist_pydantic_model(1, float, 'param_name', 'func_name')
== 1.0
)
def test_union_types():
assert (
convert_if_exist_pydantic_model(
1, Union[int, float], 'param_name', 'func_name'
)
== 1
)
assert (
convert_if_exist_pydantic_model(
1.0, Union[int, float], 'param_name', 'func_name'
)
== 1.0
)
assert (
convert_if_exist_pydantic_model(
'1.0', Union[str, float], 'param_name', 'func_name'
)
== '1.0'
)
assert (
convert_if_exist_pydantic_model(
True, Union[bool, str], 'param_name', 'func_name'
)
is True
)
# | in 3.10+
if sys.version_info >= (3, 10):
assert (
convert_if_exist_pydantic_model(1, int | float, 'param_name', 'func_name')
== 1
)
def test_nested_pydantic_model():
class SimpleModel(pydantic.BaseModel):
key1_simple: int
key2_simple: float
class ComplexModel(pydantic.BaseModel):
key1_complex: SimpleModel
key2_complex: list[SimpleModel]
key3_complex: dict[str, SimpleModel]
def foo(_: ComplexModel):
pass
annotation = inspect.signature(foo).parameters['_'].annotation
original_args = {
'key1_complex': {'key1_simple': 1, 'key2_simple': 1.0},
'key2_complex': [
{'key1_simple': 2, 'key2_simple': 2.0},
{'key1_simple': 3, 'key2_simple': 3.0},
],
'key3_complex': {
'key1_simple': {'key1_simple': 4, 'key2_simple': 4.0},
'key2_simple': {'key1_simple': 5, 'key2_simple': 5.0},
},
}
converted_args = convert_if_exist_pydantic_model(
original_args, annotation, 'param_name', 'func_name'
)
assert isinstance(converted_args, ComplexModel)
assert isinstance(converted_args.key1_complex, SimpleModel)
assert isinstance(converted_args.key2_complex[0], SimpleModel)
assert isinstance(converted_args.key2_complex[1], SimpleModel)
assert isinstance(converted_args.key3_complex['key1_simple'], SimpleModel)
assert isinstance(converted_args.key3_complex['key2_simple'], SimpleModel)
assert converted_args.model_dump() == original_args
def test_pydantic_model_in_list_union_type():
class SimpleModel(pydantic.BaseModel):
key1_simple: int
key2_simple: float
def foo(_: list[Union[int, SimpleModel]]):
pass
annotation = inspect.signature(foo).parameters['_'].annotation
original_args = [1, {'key1_simple': 1, 'key2_simple': 1.0}]
converted_args = convert_if_exist_pydantic_model(
original_args, annotation, 'param_name', 'func_name'
)
assert isinstance(converted_args, list)
assert len(converted_args) == 2
assert isinstance(converted_args[0], int)
assert isinstance(converted_args[1], SimpleModel)
assert converted_args[0] == original_args[0]
assert converted_args[1].model_dump() == original_args[1]
def test_generic_list_type():
def foo(_: Optional[list[Union[int, float]]]):
pass
annotation = inspect.signature(foo).parameters['_'].annotation
original_args = [1, 1.5]
converted_args = convert_if_exist_pydantic_model(
original_args, annotation, 'param_name', 'func_name'
)
assert isinstance(converted_args, list)
assert len(converted_args) == 2
assert isinstance(converted_args[0], int)
assert isinstance(converted_args[1], float)
assert converted_args[0] == original_args[0]
@pytest.mark.skipif(sys.version_info < (3, 10), reason='need python 3.10+')
def test_pydantic_model_with_union_and_generic_type():
class Model1(pydantic.BaseModel):
arg1: Union[int, float]
arg2: list[str]
arg3: int | float # python 3.9+
arg4: Union[list[str], str]
def foo(_: Optional[list[Union[int, Model1]]]):
pass
annotation = inspect.signature(foo).parameters['_'].annotation
original_args = [1, {'arg1': 1, 'arg2': ['a'], 'arg3': 1.0, 'arg4': '1.0'}]
converted_args = convert_if_exist_pydantic_model(
original_args, annotation, 'param_name', 'func_name'
)
assert isinstance(converted_args, list)
assert len(converted_args) == 2
assert isinstance(converted_args[0], int)
assert isinstance(converted_args[1], Model1)
assert converted_args[0] == original_args[0]
def test_unknown_pydantic_model_argument_error():
class SimpleModel(pydantic.BaseModel):
key1_simple: int
key2_simple: float
def foo(_: SimpleModel):
pass
annotation = inspect.signature(foo).parameters['_'].annotation
original_args = {'key3_simple': 1, 'key2_simple': 1.0}
with pytest.raises(errors.UnknownFunctionCallArgumentError):
convert_if_exist_pydantic_model(
original_args, annotation, 'param_name', 'func_name'
)
def test_unknown_pydantic_model_argument_error_with_union_type():
class SimpleModel1(pydantic.BaseModel):
key1_simple: int
key2_simple: float
class SimpleModel2(pydantic.BaseModel):
key3_simple: str
key4_simple: float
def foo(_: Union[SimpleModel1, SimpleModel2]):
pass
annotation = inspect.signature(foo).parameters['_'].annotation
original_args = {'key5_simple': 1, 'key4_simple': 1.0}
with pytest.raises(errors.UnknownFunctionCallArgumentError):
convert_if_exist_pydantic_model(
original_args, annotation, 'param_name', 'func_name'
)
def test_unknown_pydantic_model_argument_error_with_union_type_and_builtin_type():
class SimpleModel1(pydantic.BaseModel):
key1_simple: int
key2_simple: float
def foo(_: Union[SimpleModel1, int]):
pass
annotation = inspect.signature(foo).parameters['_'].annotation
original_args = {'key5_simple': 1, 'key4_simple': 1.0}
with pytest.raises(errors.UnknownFunctionCallArgumentError):
convert_if_exist_pydantic_model(
original_args, annotation, 'param_name', 'func_name'
)
def test_incompatible_value_and_annotation():
with pytest.raises(errors.UnknownFunctionCallArgumentError):
convert_if_exist_pydantic_model(
{'key1_simple': 1, 'key2_simple': 1.0},
int,
'param_name',
'func_name',
)
with pytest.raises(errors.UnknownFunctionCallArgumentError):
convert_if_exist_pydantic_model(
{'key1_simple': 1, 'key2_simple': 1.0},
float,
'param_name',
'func_name',
)
with pytest.raises(errors.UnknownFunctionCallArgumentError):
convert_if_exist_pydantic_model(
1,
str,
'param_name',
'func_name',
)
with pytest.raises(errors.UnknownFunctionCallArgumentError):
convert_if_exist_pydantic_model(
1.0,
str,
'param_name',
'func_name',
)
with pytest.raises(errors.UnknownFunctionCallArgumentError):
convert_if_exist_pydantic_model(
True,
str,
'param_name',
'func_name',
)
with pytest.raises(errors.UnknownFunctionCallArgumentError):
convert_if_exist_pydantic_model(
[1],
str,
'param_name',
'func_name',
)
with pytest.raises(errors.UnknownFunctionCallArgumentError):
convert_if_exist_pydantic_model(
{'key1': 1},
str,
'param_name',
'func_name',
)
with pytest.raises(errors.UnknownFunctionCallArgumentError):
convert_if_exist_pydantic_model(
{'key1': 1, 'key2': 2},
str,
'param_name',
'func_name',
)

View File

@@ -0,0 +1,63 @@
# 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.
#
"""Tests for convert_number_values_for_function_call_args."""
from ..._extra_utils import convert_number_values_for_function_call_args
def test_integer_value():
assert convert_number_values_for_function_call_args(1) == 1
def test_float_value():
assert convert_number_values_for_function_call_args(1.0) == 1
def test_string_value():
assert convert_number_values_for_function_call_args('1.0') == '1.0'
def test_boolean_value():
assert convert_number_values_for_function_call_args(True) is True
def test_none_value():
assert convert_number_values_for_function_call_args(None) is None
def test_float_value_with_decimal():
assert convert_number_values_for_function_call_args(1.1) == 1.1
def test_dict_value():
assert convert_number_values_for_function_call_args(
{'key1': 1.0, 'key2': 1.1}
) == {'key1': 1, 'key2': 1.1}
def test_list_value():
assert convert_number_values_for_function_call_args([1.0, 1.1, 1.2]) == [
1,
1.1,
1.2,
]
def test_nested_value():
assert convert_number_values_for_function_call_args(
{'key1': 1.0, 'key2': {'key3': 1.0, 'key4': [1.2, 2.0]}}
) == {'key1': 1, 'key2': {'key3': 1, 'key4': [1.2, 2]}}

View File

@@ -0,0 +1,268 @@
# 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.
#
"""Tests for find_afc_incompatible_tool_indexes."""
from typing import Any
import pytest
from ... import types
from ..._extra_utils import find_afc_incompatible_tool_indexes
try:
from mcp import types as mcp_types
from mcp import ClientSession as McpClientSession
from ..._adapters import McpToGenAiToolAdapter
except ImportError as e:
import sys
if sys.version_info < (3, 10):
raise ImportError(
'MCP Tool requires Python 3.10 or above. Please upgrade your Python'
' version.'
) from e
else:
raise e
def get_weather_tool(city: str) -> str:
"""Get the weather in a city."""
return f'The weather in {city} is sunny and 100 degrees.'
class MockMcpClientSession(McpClientSession):
def __init__(self):
self._read_stream = None
self._write_stream = None
async def call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
return mcp_types.CallToolResult(
content=[mcp_types.TextContent(type='text', text='1.01')]
)
mcp_to_genai_tool_adapter = McpToGenAiToolAdapter(
session=MockMcpClientSession(),
list_tools_result=mcp_types.ListToolsResult(tools=[]),
)
def test_no_config_returns_empty_list():
"""Verifies that an empty list is returned if the input config is None.
"""
result = find_afc_incompatible_tool_indexes(config=None)
assert result == []
def test_config_with_no_tools_returns_empty_list():
"""Verifies an empty list is returned if the config has no 'tools' attribute.
"""
result = find_afc_incompatible_tool_indexes(
config=types.GenerateContentConfig()
)
assert result == []
def test_empty_tools_list_returns_empty_list():
"""Verifies that an empty list is returned if the 'tools' list is empty."""
result = find_afc_incompatible_tool_indexes(
config=types.GenerateContentConfig(tools=[])
)
assert result == []
def test_all_compatible_tools_returns_empty_list_with_empty_fd():
"""Verifies that an empty list is returned when all tools are compatible.
A tool is compatible if it's not a `types.Tool` or if its
`function_declarations` attribute is empty or None from config.
"""
result = find_afc_incompatible_tool_indexes(
config=types.GenerateContentConfig(
tools=[
types.Tool(
google_search_retrieval=types.GoogleSearchRetrieval()
),
types.Tool(retrieval=types.Retrieval()),
types.Tool(google_search=types.GoogleSearch()),
types.Tool(code_execution=types.ToolCodeExecution()),
types.Tool(google_maps=types.GoogleMaps()),
types.Tool(url_context=types.UrlContext()),
types.Tool(computer_use=types.ComputerUse()),
types.Tool(code_execution=types.ToolCodeExecution()),
types.Tool(function_declarations=[]),
mcp_types.Tool(
name='get_weather',
description='Get the weather in a city.',
inputSchema={
'type': 'object',
'properties': {'location': {'type': 'string'}},
},
),
get_weather_tool,
mcp_to_genai_tool_adapter,
]
)
)
assert result == []
def test_all_compatible_tools_returns_empty_list_with_none_fd():
"""Verifies that an empty list is returned when all tools are compatible.
A tool is compatible if it's not a `types.Tool` or if its
`function_declarations` attribute is empty or None from config.
"""
result = find_afc_incompatible_tool_indexes(
config=types.GenerateContentConfig(
tools=[
types.Tool(
google_search_retrieval=types.GoogleSearchRetrieval()
),
types.Tool(retrieval=types.Retrieval()),
types.Tool(google_search=types.GoogleSearch()),
types.Tool(code_execution=types.ToolCodeExecution()),
types.Tool(google_maps=types.GoogleMaps()),
types.Tool(url_context=types.UrlContext()),
types.Tool(computer_use=types.ComputerUse()),
types.Tool(code_execution=types.ToolCodeExecution()),
types.Tool(function_declarations=None),
mcp_types.Tool(
name='get_weather',
description='Get the weather in a city.',
inputSchema={
'type': 'object',
'properties': {'location': {'type': 'string'}},
},
),
get_weather_tool,
mcp_to_genai_tool_adapter,
]
)
)
assert result == []
def test_all_compatible_tools_returns_empty_list():
"""Verifies that an empty list is returned when all tools are compatible.
A tool is compatible if it's not a `types.Tool` or if its
`function_declarations` attribute is empty or None from config.
"""
result = find_afc_incompatible_tool_indexes(
config=types.GenerateContentConfig(
tools=[
types.Tool(
google_search_retrieval=types.GoogleSearchRetrieval()
),
types.Tool(retrieval=types.Retrieval()),
types.Tool(google_search=types.GoogleSearch()),
types.Tool(code_execution=types.ToolCodeExecution()),
types.Tool(google_maps=types.GoogleMaps()),
types.Tool(url_context=types.UrlContext()),
types.Tool(computer_use=types.ComputerUse()),
types.Tool(code_execution=types.ToolCodeExecution()),
types.Tool(function_declarations=[]),
mcp_types.Tool(
name='get_weather',
description='Get the weather in a city.',
inputSchema={
'type': 'object',
'properties': {'location': {'type': 'string'}},
},
),
get_weather_tool,
mcp_to_genai_tool_adapter,
]
)
)
assert result == []
def test_single_incompatible_tool():
"""Verifies that the correct index is returned for a single incompatible
tool.
"""
result = find_afc_incompatible_tool_indexes(
config=types.GenerateContentConfig(
tools=[
types.Tool(
google_search_retrieval=types.GoogleSearchRetrieval()
),
types.Tool(retrieval=types.Retrieval()),
types.Tool(
function_declarations=[
types.FunctionDeclaration(name='test_function')
]
),
get_weather_tool,
mcp_to_genai_tool_adapter,
]
)
)
assert result == [2]
def test_multiple_incompatible_tools():
"""Verifies correct indexes are returned for multiple incompatible tools."""
result = find_afc_incompatible_tool_indexes(
config=types.GenerateContentConfig(
tools=[
types.Tool(
google_search_retrieval=types.GoogleSearchRetrieval()
),
types.Tool(retrieval=types.Retrieval()),
types.Tool(
function_declarations=[
types.FunctionDeclaration(name='test_function')
]
),
types.Tool(computer_use=types.ComputerUse()),
types.Tool(code_execution=types.ToolCodeExecution()),
types.Tool(
function_declarations=[
types.FunctionDeclaration(name='test_function_2')
]
),
get_weather_tool,
mcp_to_genai_tool_adapter,
]
)
)
assert result == [2, 5]
def test_mcp_tool_incompatible():
"""Verifies correct indexes are returned for multiple incompatible tools."""
result = find_afc_incompatible_tool_indexes(
config=types.GenerateContentConfig(
tools=[
types.Tool(
google_search_retrieval=types.GoogleSearchRetrieval()
),
types.Tool(retrieval=types.Retrieval()),
types.Tool(
function_declarations=[
types.FunctionDeclaration(name='test_function')
]
),
types.Tool(code_execution=types.ToolCodeExecution()),
get_weather_tool,
mcp_to_genai_tool_adapter,
types.Tool(
mcp_servers=[types.McpServer(name='test_mcp_server')]
),
]
)
)
assert result == [2, 6]

View File

@@ -0,0 +1,530 @@
# 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.
#
from unittest import mock
import pytest
from ... import _api_client
from ... import _extra_utils
from ... import client
from ... import models
from ... import types
TEST_NO_AFC_PART = types.Part(
text=(
'Okay, here is the weather in San Francisco'
' as of approximately 8:10 pm PST on May'
' 10, 2023. Please note that the weather'
' can change rapidly.'
)
)
TEST_FUNCTION_CALL_PART = types.Part(
function_call=types.FunctionCall(
name='get_current_weather',
args={'location': 'San Francisco'},
)
)
TEST_FUNCTION_RESPONSE_PART = types.Part(
function_response=types.FunctionResponse(
name='get_current_weather',
response={'result': 'sunny'},
)
)
TEST_AFC_TEXT_PART = types.Part(text='San Francisco weather is sunny.')
TEST_NO_AFC_CONTENT = types.Content(
parts=[TEST_NO_AFC_PART],
role='model',
)
TEST_FUNCTION_CALL_CONTENT = types.Content(
parts=[TEST_FUNCTION_CALL_PART],
role='model',
)
TEST_FUNCTION_RESPONSE_CONTENT = types.Content(
parts=[TEST_FUNCTION_RESPONSE_PART],
role='user',
)
TEST_AFC_TEXT_CONTENT = types.Content(
parts=[TEST_AFC_TEXT_PART],
role='model',
)
TEST_AFC_HISTORY = [
types.Content(
parts=[types.Part(text='what is the weather in San Francisco?')],
role='user',
),
TEST_FUNCTION_CALL_CONTENT,
TEST_FUNCTION_RESPONSE_CONTENT,
]
def get_current_weather(location: str) -> str:
"""Returns the current weather.
Args:
location: The location of a city and state, e.g. "San Francisco, CA".
"""
return 'windy'
async def get_current_weather_async(location: str) -> str:
"""Returns the current weather.
Args:
location: The location of a city and state, e.g. "San Francisco, CA".
"""
return 'windy'
def get_aqi_from_city(location: str) -> str:
"""Returns the aqi index of a city.
Args:
location: The city and State, e.g. San Francisco, CA.
"""
return None
@pytest.fixture
def mock_api_client(vertexai=False):
api_client = mock.MagicMock(spec=client.ApiClient)
api_client.api_key = 'TEST_API_KEY'
api_client._host = lambda: 'test_host'
api_client._http_options = {'headers': {}} # Ensure headers exist
api_client.vertexai = vertexai
return api_client
@pytest.fixture
def mock_get_function_response_parts_none():
with mock.patch.object(
_extra_utils,
'get_function_response_parts',
) as mock_get_function_response_parts_none:
mock_get_function_response_parts_none.return_value = None
yield mock_get_function_response_parts_none
@pytest.fixture
def mock_get_function_response_parts_none_async():
with mock.patch.object(
_extra_utils,
'get_function_response_parts_async',
) as mock_get_function_response_parts_none_async:
mock_get_function_response_parts_none_async.return_value = None
yield mock_get_function_response_parts_none_async
@pytest.fixture
def mock_get_function_response_parts() -> list[types.Part]:
with mock.patch.object(
_extra_utils, 'get_function_response_parts'
) as mock_get_function_response_parts:
mock_get_function_response_parts.side_effect = [
[TEST_FUNCTION_RESPONSE_PART],
[], # Breaks when the function response is not returned.
]
yield mock_get_function_response_parts
@pytest.fixture
def mock_get_function_response_parts_async() -> list[types.Part]:
with mock.patch.object(
_extra_utils, 'get_function_response_parts_async'
) as mock_get_function_response_parts_async:
mock_get_function_response_parts_async.side_effect = [
[TEST_FUNCTION_RESPONSE_PART],
[], # Breaks when the function response is not returned.
]
yield mock_get_function_response_parts_async
@pytest.fixture
def mock_generate_content_stream_no_afc():
with mock.patch.object(
models.Models, '_generate_content_stream'
) as mock_stream_no_afc:
mock_stream_no_afc.return_value = [
types.GenerateContentResponse(
candidates=[types.Candidate(content=TEST_NO_AFC_CONTENT)]
)
]
yield mock_stream_no_afc
@pytest.fixture
def mock_generate_content_stream_with_afc():
with mock.patch.object(
models.Models, '_generate_content_stream'
) as mock_stream_with_afc:
mock_stream_with_afc.side_effect = [
[
types.GenerateContentResponse(
candidates=[types.Candidate(content=TEST_FUNCTION_CALL_CONTENT)]
)
],
[
types.GenerateContentResponse(
candidates=[types.Candidate(content=TEST_AFC_TEXT_CONTENT)]
)
],
]
yield mock_stream_with_afc
@pytest.fixture
def mock_generate_content_stream_no_afc_async():
with mock.patch.object(
models.AsyncModels, '_generate_content_stream'
) as mock_stream_no_afc:
async def async_generator():
yield types.GenerateContentResponse(
candidates=[types.Candidate(content=TEST_NO_AFC_CONTENT)]
)
mock_stream_no_afc.return_value = async_generator()
yield mock_stream_no_afc
@pytest.fixture
def mock_generate_content_stream_with_afc_async():
with mock.patch.object(
models.AsyncModels, '_generate_content_stream'
) as mock_stream_with_afc:
async def async_generator_1():
yield types.GenerateContentResponse(
candidates=[types.Candidate(content=TEST_FUNCTION_CALL_CONTENT)]
)
async def async_generator_2():
yield types.GenerateContentResponse(
candidates=[types.Candidate(content=TEST_AFC_TEXT_CONTENT)]
)
mock_stream_with_afc.side_effect = [
async_generator_1(),
async_generator_2(),
]
yield mock_stream_with_afc
def test_generate_content_stream_no_function_map(
mock_generate_content_stream_no_afc,
mock_get_function_response_parts_none,
):
"""Test when no function tools are provided.
Expected to answer past weather.
"""
models_instance = models.Models(api_client_=mock_api_client)
stream = models_instance.generate_content_stream(
model='test_model', contents='what is the weather in San Francisco?'
)
for chunk in stream:
assert chunk.text == TEST_NO_AFC_PART.text
assert mock_generate_content_stream_no_afc.call_count == 1
assert mock_get_function_response_parts_none.call_count == 0
def test_generate_content_stream_afc_disabled(
mock_generate_content_stream_with_afc,
mock_get_function_response_parts_none,
):
"""Test when function tools are provided but AFC is disabled.
Expected to respond with function call.
"""
models_instance = models.Models(api_client_=mock_api_client)
stream = models_instance.generate_content_stream(
model='test_model',
contents='what is the weather in San Francisco?',
config=types.GenerateContentConfig(
tools=[get_current_weather],
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=True
),
),
)
for chunk in stream:
# Work as manual function calling.
assert chunk.candidates[0].content.parts[0].function_call
assert mock_generate_content_stream_with_afc.call_count == 1
assert mock_get_function_response_parts_none.call_count == 0
def test_generate_content_stream_no_function_response(
mock_generate_content_stream_no_afc,
mock_get_function_response_parts_none,
):
"""Test when function tools are provided and function responses are not returned.
Expected to answer past weather.
"""
models_instance = models.Models(api_client_=mock_api_client)
config = types.GenerateContentConfig(tools=[get_aqi_from_city])
stream = models_instance.generate_content_stream(
model='test_model',
contents='what is the weather in San Francisco?',
config=config,
)
for chunk in stream:
assert chunk.text == TEST_NO_AFC_PART.text
assert mock_generate_content_stream_no_afc.call_count == 1
assert mock_get_function_response_parts_none.call_count == 1
def test_generate_content_stream_with_function_tools_used(
mock_generate_content_stream_with_afc,
mock_get_function_response_parts,
):
"""Test when function tools are provided and function responses are returned.
Expected to answer weather based on function response.
"""
models_instance = models.Models(api_client_=mock_api_client)
config = types.GenerateContentConfig(tools=[get_current_weather])
stream = models_instance.generate_content_stream(
model='test_model',
contents='what is the weather in San Francisco?',
config=config,
)
chunk = None
for chunk in stream:
assert chunk.text == TEST_AFC_TEXT_PART.text
assert mock_generate_content_stream_with_afc.call_count == 2
assert mock_get_function_response_parts.call_count == 2
assert chunk is not None
for i in range(len(chunk.automatic_function_calling_history)):
assert chunk.automatic_function_calling_history[i].model_dump(
exclude_none=True
) == TEST_AFC_HISTORY[i].model_dump(exclude_none=True)
def test_generate_content_stream_with_thought_summaries(
mock_generate_content_stream_with_afc,
mock_get_function_response_parts,
):
"""Test when function tools are provided and thought summaries are enabled.
Expected to answer weather based on function response.
"""
models_instance = models.Models(api_client_=mock_api_client)
config = types.GenerateContentConfig(
tools=[get_current_weather],
thinking_config=types.ThinkingConfig(include_thoughts=True),
)
stream = models_instance.generate_content_stream(
model='test_model',
contents='what is the weather in San Francisco?',
config=config,
)
chunk = None
for chunk in stream:
assert chunk.text == TEST_AFC_TEXT_PART.text
assert mock_generate_content_stream_with_afc.call_count == 2
assert mock_get_function_response_parts.call_count == 2
assert chunk is not None
for i in range(len(chunk.automatic_function_calling_history)):
assert chunk.automatic_function_calling_history[i].model_dump(
exclude_none=True
) == TEST_AFC_HISTORY[i].model_dump(exclude_none=True)
@pytest.mark.asyncio
async def test_generate_content_stream_no_function_map_async(
mock_generate_content_stream_no_afc,
mock_get_function_response_parts_none,
):
"""Test when no function tools are provided.
Expected to answer past weather.
"""
models_instance = models.Models(api_client_=mock_api_client)
stream = models_instance.generate_content_stream(
model='test_model', contents='what is the weather in San Francisco?'
)
for chunk in stream:
assert chunk.text == TEST_NO_AFC_PART.text
assert mock_generate_content_stream_no_afc.call_count == 1
assert mock_get_function_response_parts_none.call_count == 0
@pytest.mark.asyncio
async def test_generate_content_stream_afc_disabled_async(
mock_generate_content_stream_with_afc_async,
mock_get_function_response_parts_none,
):
"""Test when function tools are provided but AFC is disabled.
Expected to respond with function call.
"""
models_instance = models.AsyncModels(api_client_=mock_api_client)
stream = await models_instance.generate_content_stream(
model='test_model',
contents='what is the weather in San Francisco?',
config=types.GenerateContentConfig(
tools=[get_current_weather],
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=True
),
),
)
async for chunk in stream:
# Work as manual function calling.
assert chunk.candidates[0].content.parts[0].function_call
assert mock_generate_content_stream_with_afc_async.call_count == 1
assert mock_get_function_response_parts_none.call_count == 0
@pytest.mark.asyncio
async def test_generate_content_stream_no_function_response_async(
mock_generate_content_stream_no_afc_async,
mock_get_function_response_parts_none_async,
):
"""Test when function tools are provided and function responses are not returned.
Expected to answer past weather.
"""
models_instance = models.AsyncModels(api_client_=mock_api_client)
config = types.GenerateContentConfig(tools=[get_aqi_from_city])
stream = await models_instance.generate_content_stream(
model='test_model',
contents='what is the weather in San Francisco?',
config=config,
)
async for chunk in stream:
assert chunk.text == TEST_NO_AFC_PART.text
assert mock_generate_content_stream_no_afc_async.call_count == 1
assert mock_get_function_response_parts_none_async.call_count == 1
@pytest.mark.asyncio
async def test_generate_content_stream_with_function_tools_used_async(
mock_generate_content_stream_with_afc_async,
mock_get_function_response_parts_async,
):
"""Test when function tools are provided and function responses are returned.
Expected to answer weather based on function response.
"""
models_instance = models.AsyncModels(api_client_=mock_api_client)
config = types.GenerateContentConfig(tools=[get_current_weather])
stream = await models_instance.generate_content_stream(
model='test_model',
contents='what is the weather in San Francisco?',
config=config,
)
chunk = None
async for chunk in stream:
assert chunk.text == TEST_AFC_TEXT_PART.text
assert mock_generate_content_stream_with_afc_async.call_count == 2
assert mock_get_function_response_parts_async.call_count == 2
assert chunk is not None
for i in range(len(chunk.automatic_function_calling_history)):
assert chunk.automatic_function_calling_history[i].model_dump(
exclude_none=True
) == TEST_AFC_HISTORY[i].model_dump(exclude_none=True)
@pytest.mark.asyncio
async def test_generate_content_stream_with_function_async_function_used_async(
mock_generate_content_stream_with_afc_async,
mock_get_function_response_parts_async,
):
"""Test when function tools are provided and function responses are returned.
Expected to answer weather based on function response.
"""
models_instance = models.AsyncModels(api_client_=mock_api_client)
config = types.GenerateContentConfig(tools=[get_current_weather_async])
stream = await models_instance.generate_content_stream(
model='test_model',
contents='what is the weather in San Francisco?',
config=config,
)
chunk = None
async for chunk in stream:
assert chunk.text == TEST_AFC_TEXT_PART.text
assert mock_generate_content_stream_with_afc_async.call_count == 2
assert mock_get_function_response_parts_async.call_count == 2
assert chunk is not None
for i in range(len(chunk.automatic_function_calling_history)):
assert chunk.automatic_function_calling_history[i].model_dump(
exclude_none=True
) == TEST_AFC_HISTORY[i].model_dump(exclude_none=True)
@pytest.mark.asyncio
async def test_generate_content_stream_with_thought_summaries_async(
mock_generate_content_stream_with_afc_async,
mock_get_function_response_parts_async,
):
"""Test when function tools are provided and thought summaries are enabled.
Expected to answer weather based on function response.
"""
models_instance = models.AsyncModels(api_client_=mock_api_client)
config = types.GenerateContentConfig(
tools=[get_current_weather],
thinking_config=types.ThinkingConfig(include_thoughts=True),
)
stream = await models_instance.generate_content_stream(
model='test_model',
contents='what is the weather in San Francisco?',
config=config,
)
chunk = None
async for chunk in stream:
assert chunk.text == TEST_AFC_TEXT_PART.text
assert mock_generate_content_stream_with_afc_async.call_count == 2
assert mock_get_function_response_parts_async.call_count == 2
assert chunk is not None
for i in range(len(chunk.automatic_function_calling_history)):
assert chunk.automatic_function_calling_history[i].model_dump(
exclude_none=True
) == TEST_AFC_HISTORY[i].model_dump(exclude_none=True)

View File

@@ -0,0 +1,77 @@
# 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 pytest
from ... import types
from .. import pytest_helper
def get_current_weather(location: str) -> str:
"""Returns the current weather.
Args:
location: The location of a city and state, e.g. "San Francisco, CA".
"""
return 'windy'
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='models.generate_content',
)
pytest_plugins = ('pytest_asyncio',)
def test_generate_content_stream_with_function_and_thought_summaries(client):
"""Test when function tools are provided and thought summaries are enabled.
Expected to answer weather based on function response.
"""
config = types.GenerateContentConfig(
tools=[get_current_weather],
thinking_config=types.ThinkingConfig(include_thoughts=True),
)
stream = client.models.generate_content_stream(
model='gemini-2.5-flash',
contents='what is the weather in San Francisco, CA?',
config=config,
)
chunk = None
for chunk in stream:
assert chunk is not None
@pytest.mark.asyncio
async def test_generate_content_stream_with_function_and_thought_summaries_async(
client,
):
"""Test when function tools are provided and thought summaries are enabled.
Expected to answer weather based on function response.
"""
config = types.GenerateContentConfig(
tools=[get_current_weather],
thinking_config=types.ThinkingConfig(include_thoughts=True),
)
stream = await client.aio.models.generate_content_stream(
model='gemini-2.5-flash',
contents='what is the weather in San Francisco, CA?',
config=config,
)
chunk = None
async for chunk in stream:
assert chunk is not None

View File

@@ -0,0 +1,176 @@
# 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.
#
"""Tests for get_function_map."""
import typing
from typing import Any
import pytest
from ..._extra_utils import get_function_map
from ...errors import UnsupportedFunctionError
from ...types import GenerateContentConfig
_is_mcp_imported = False
if typing.TYPE_CHECKING:
from mcp import types as mcp_types
from mcp import ClientSession as McpClientSession
from ..._adapters import McpToGenAiToolAdapter
_is_mcp_imported = True
else:
McpClientSession: typing.Type = Any
McpToGenAiToolAdapter: typing.Type = Any
try:
from mcp import types as mcp_types
from mcp import ClientSession as McpClientSession
from ..._adapters import McpToGenAiToolAdapter
_is_mcp_imported = True
except ImportError:
McpClientSession = None
McpToGenAiToolAdapter = None
def test_coroutine_function():
async def func_under_test():
pass
config = GenerateContentConfig(tools=[func_under_test])
with pytest.raises(UnsupportedFunctionError):
get_function_map(config)
def test_empty_config():
config = {}
assert get_function_map(config) == {}
def test_empty_tools():
config = GenerateContentConfig(top_p=0.5)
assert get_function_map(config) == {}
def test_valid_function():
def func_under_test():
pass
config = GenerateContentConfig(tools=[func_under_test])
assert get_function_map(config) == {'func_under_test': func_under_test}
def test_mcp_tool_raises_error():
if not _is_mcp_imported:
return
session = McpClientSession(read_stream=None, write_stream=None)
config = GenerateContentConfig(tools=[session])
mcp_to_genai_tool_adapters = {'tool': McpToGenAiToolAdapter(session, [])}
with pytest.raises(UnsupportedFunctionError):
get_function_map(
config, mcp_to_genai_tool_adapters, is_caller_method_async=False
)
@pytest.mark.asyncio
async def test_mcp_tool():
if not _is_mcp_imported:
return
class MockMcpClientSession(McpClientSession):
def __init__(self):
self._read_stream = None
self._write_stream = None
async def list_tools(self):
return mcp_types.ListToolsResult(
tools=[
mcp_types.Tool(
name='tool',
description='tool-description',
inputSchema={
'type': 'OBJECT',
'properties': {
'key1': {
'type': 'STRING',
},
'key2': {
'type': 'NUMBER',
},
},
},
)
]
)
session = MockMcpClientSession()
config = GenerateContentConfig(tools=[session])
mcp_to_genai_tool_adapters = {
'tool': McpToGenAiToolAdapter(session, [await session.list_tools()]),
}
result = get_function_map(
config, mcp_to_genai_tool_adapters, is_caller_method_async=True
)
assert isinstance(result['tool'], McpToGenAiToolAdapter)
@pytest.mark.asyncio
async def test_duplicate_mcp_tool_raises_error():
if not _is_mcp_imported:
return
class MockMcpClientSession(McpClientSession):
def __init__(self):
self._read_stream = None
self._write_stream = None
async def list_tools(self):
return mcp_types.ListToolsResult(
tools=[
mcp_types.Tool(
name='tool',
description='tool-description',
inputSchema={
'type': 'OBJECT',
'properties': {
'key1': {
'type': 'STRING',
},
'key2': {
'type': 'NUMBER',
},
},
},
)
]
)
def tool():
pass
session = MockMcpClientSession()
config = GenerateContentConfig(tools=[tool, session])
mcp_to_genai_tool_adapters = {
'tool': McpToGenAiToolAdapter(session, [await session.list_tools()]),
}
with pytest.raises(ValueError):
get_function_map(
config, mcp_to_genai_tool_adapters, is_caller_method_async=True
)

View File

@@ -0,0 +1,277 @@
# 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.
#
"""Tests for get_function_response_parts."""
import typing
from typing import Any
import pytest
from ..._extra_utils import get_function_response_parts, get_function_response_parts_async
from ...errors import UnsupportedFunctionError
from ...types import Candidate
from ...types import Content
from ...types import FunctionCall
from ...types import FunctionResponse
from ...types import GenerateContentResponse
from ...types import Part
_is_mcp_imported = False
if typing.TYPE_CHECKING:
from mcp import types as mcp_types
from mcp import ClientSession as McpClientSession
from ..._adapters import McpToGenAiToolAdapter
_is_mcp_imported = True
else:
McpClientSession: typing.Type = Any
McpToGenAiToolAdapter: typing.Type = Any
try:
from mcp import types as mcp_types
from mcp import ClientSession as McpClientSession
from ..._adapters import McpToGenAiToolAdapter
_is_mcp_imported = True
except ImportError:
McpClientSession = None
McpToGenAiToolAdapter = None
def test_integer_value():
def func_under_test(a: int) -> int:
return a + 1
response = GenerateContentResponse(
candidates=[
Candidate(
content=Content(
parts=[
Part(
function_call=FunctionCall(
name='func_under_test',
args={'a': 1},
)
)
]
)
)
]
)
function_map = {'func_under_test': func_under_test}
expected_parts = [
Part(
function_response=FunctionResponse(
name='func_under_test',
response={'result': 2},
)
)
]
actual_parts = get_function_response_parts(response, function_map)
for actual_part, expected_part in zip(actual_parts, expected_parts):
assert actual_part.model_dump_json(
exclude_none=True
) == expected_part.model_dump_json(exclude_none=True)
def test_float_value():
def func_under_test(a: float) -> float:
return a + 1.0
response = GenerateContentResponse(
candidates=[
Candidate(
content=Content(
parts=[
Part(
function_call=FunctionCall(
name='func_under_test',
args={'a': 1.0},
)
)
]
)
)
]
)
function_map = {'func_under_test': func_under_test}
expected_parts = [
Part(
function_response=FunctionResponse(
name='func_under_test',
response={'result': 2.0},
)
)
]
actual_parts = get_function_response_parts(response, function_map)
for actual_part, expected_part in zip(actual_parts, expected_parts):
assert actual_part.model_dump_json(
exclude_none=True
) == expected_part.model_dump_json(exclude_none=True)
def test_string_value():
def func_under_test(a: str) -> str:
return a + '1'
response = GenerateContentResponse(
candidates=[
Candidate(
content=Content(
parts=[
Part(
function_call=FunctionCall(
name='func_under_test',
args={'a': '1.0'},
)
)
]
)
)
]
)
function_map = {'func_under_test': func_under_test}
expected_parts = [
Part(
function_response=FunctionResponse(
name='func_under_test',
response={'result': '1.01'},
)
)
]
actual_parts = get_function_response_parts(response, function_map)
for actual_part, expected_part in zip(actual_parts, expected_parts):
assert actual_part.model_dump_json(
exclude_none=True
) == expected_part.model_dump_json(exclude_none=True)
@pytest.mark.asyncio
async def test_mcp_tool():
if not _is_mcp_imported:
return
class MockMcpClientSession(McpClientSession):
def __init__(self):
self._read_stream = None
self._write_stream = None
async def call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
return mcp_types.CallToolResult(
content=[mcp_types.TextContent(type='text', text='1.01')]
)
mcp_to_genai_tool_adapter = McpToGenAiToolAdapter(
session=MockMcpClientSession(),
list_tools_result=mcp_types.ListToolsResult(tools=[]),
)
response = GenerateContentResponse(
candidates=[
Candidate(
content=Content(
parts=[
Part(
function_call=FunctionCall(
name='tool',
args={'key1': 'value1', 'key2': 1},
)
)
]
)
)
]
)
function_map = {'tool': mcp_to_genai_tool_adapter}
expected_parts = [
Part(
function_response=FunctionResponse(
name='tool',
response={
'result': {
'content': [{'type': 'text', 'text': '1.01'}],
'isError': False,
}
},
)
)
]
actual_parts = await get_function_response_parts_async(response, function_map)
for actual_part, expected_part in zip(actual_parts, expected_parts):
assert actual_part.model_dump_json(
exclude_none=True
) == expected_part.model_dump_json(exclude_none=True)
@pytest.mark.asyncio
async def test_mcp_tool_error():
if not _is_mcp_imported:
return
class MockMcpClientSession(McpClientSession):
def __init__(self):
self._read_stream = None
self._write_stream = None
async def call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
return mcp_types.CallToolResult(
content=[mcp_types.TextContent(type='text', text='Internal error')],
isError=True,
)
mcp_to_genai_tool_adapter = McpToGenAiToolAdapter(
session=MockMcpClientSession(),
list_tools_result=mcp_types.ListToolsResult(tools=[]),
)
response = GenerateContentResponse(
candidates=[
Candidate(
content=Content(
parts=[
Part(
function_call=FunctionCall(
name='tool',
args={'key1': 'value1', 'key2': 1},
)
)
]
)
)
]
)
function_map = {'tool': mcp_to_genai_tool_adapter}
expected_parts = [
Part(
function_response=FunctionResponse(
name='tool',
response={
'error': {
'content': [{'type': 'text', 'text': 'Internal error'}],
'isError': True,
}
},
)
)
]
actual_parts = await get_function_response_parts_async(response, function_map)
for actual_part, expected_part in zip(actual_parts, expected_parts):
assert actual_part.model_dump_json(
exclude_none=True
) == expected_part.model_dump_json(exclude_none=True)

View File

@@ -0,0 +1,130 @@
# 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.
#
"""Tests for get_max_remote_calls_for_afc."""
from ... import types
from ..._extra_utils import get_max_remote_calls_afc
import pytest
def test_config_is_none():
assert get_max_remote_calls_afc(None) == 10
def test_afc_unset_max_unset():
assert get_max_remote_calls_afc(types.GenerateContentConfig()) == 10
def test_afc_unset_max_set():
assert (
get_max_remote_calls_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
maximum_remote_calls=20,
),
)
)
== 20
)
def test_afc_disabled_max_unset():
with pytest.raises(ValueError):
get_max_remote_calls_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=True,
),
)
)
def test_afc_disabled_max_set():
with pytest.raises(ValueError):
get_max_remote_calls_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=True,
maximum_remote_calls=20,
),
)
)
def test_afc_d_max_unset():
assert (
get_max_remote_calls_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=False,
),
)
)
== 10
)
def test_afc_d_max_set():
assert (
get_max_remote_calls_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=False,
maximum_remote_calls=5,
),
)
)
== 5
)
def test_afc_enabled_max_set_to_zero():
with pytest.raises(ValueError):
get_max_remote_calls_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=False,
maximum_remote_calls=0,
),
)
)
def test_afc_enabled_max_set_to_negative():
with pytest.raises(ValueError):
get_max_remote_calls_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=False,
maximum_remote_calls=-1,
),
)
)
def test_afc_enabled_max_set_to_float():
assert (
get_max_remote_calls_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=False,
maximum_remote_calls=5.0,
),
)
)
== 5
)

View File

@@ -0,0 +1,241 @@
# 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.
#
"""Test invoke_function_from_dict_args."""
from typing import Union
import pydantic
import pytest
import sys
from ... import errors
from ..._extra_utils import invoke_function_from_dict_args
def test_builtin_primitive_types():
def func_under_test(x: int, y: float, z: str, w: bool):
return {
'x': x + 1,
'y': y + 1.0,
'z': z + '1.0',
'w': not w,
}
original_args = {
'x': 1,
'y': 1.0,
'z': '1.0',
'w': True,
}
expected_response = {
'x': 2,
'y': 2.0,
'z': '1.01.0',
'w': False,
}
actual_response = invoke_function_from_dict_args(
original_args, func_under_test
)
assert actual_response == expected_response
@pytest.mark.skipif(
sys.version_info < (3, 10),
reason='| is only supported in Python 3.9 and above.',
)
def test_builtin_compound_types():
def func_under_test(x: list[int], y: dict[str, float]):
return {
'new_x': x + [3],
'new_y': y | {'key3': 3.0},
}
original_args = {
'x': [1, 2],
'y': {'key1': 1.0, 'key2': 2.0},
}
expected_response = {
'new_x': [1, 2, 3],
'new_y': {'key1': 1.0, 'key2': 2.0, 'key3': 3.0},
}
actual_response = invoke_function_from_dict_args(
original_args, func_under_test
)
assert actual_response == expected_response
def test_nested_pydantic_model():
class SimpleModel(pydantic.BaseModel):
key1_simple: int
key2_simple: float
class ComplexModel(pydantic.BaseModel):
key1_complex: SimpleModel
key2_complex: list[SimpleModel]
key3_complex: dict[str, SimpleModel]
def func_under_test(x: ComplexModel):
y = x.model_copy()
y.key1_complex.key1_simple += 1
y.key1_complex.key2_simple += 1.0
for simple_model in y.key2_complex:
simple_model.key1_simple += 1
simple_model.key2_simple += 1.0
for simple_model in y.key3_complex.values():
simple_model.key1_simple += 1
simple_model.key2_simple += 1.0
return y
original_args = {
'x': {
'key1_complex': {'key1_simple': 1, 'key2_simple': 1.0},
'key2_complex': [
{'key1_simple': 2, 'key2_simple': 2.0},
{'key1_simple': 3, 'key2_simple': 3.0},
],
'key3_complex': {
'key1_simple': {'key1_simple': 4, 'key2_simple': 4.0},
'key2_simple': {'key1_simple': 5, 'key2_simple': 5.0},
},
}
}
expected_response = {
'key1_complex': {'key1_simple': 2, 'key2_simple': 2.0},
'key2_complex': [
{'key1_simple': 3, 'key2_simple': 3.0},
{'key1_simple': 4, 'key2_simple': 4.0},
],
'key3_complex': {
'key1_simple': {'key1_simple': 5, 'key2_simple': 5.0},
'key2_simple': {'key1_simple': 6, 'key2_simple': 6.0},
},
}
actual_response = invoke_function_from_dict_args(
original_args, func_under_test
)
assert isinstance(actual_response, ComplexModel)
assert isinstance(actual_response.key1_complex, SimpleModel)
assert isinstance(actual_response.key2_complex[0], SimpleModel)
assert isinstance(actual_response.key2_complex[1], SimpleModel)
assert isinstance(actual_response.key3_complex['key1_simple'], SimpleModel)
assert isinstance(actual_response.key3_complex['key2_simple'], SimpleModel)
assert actual_response.model_dump() == expected_response
def test_pydantic_model_in_list_union_type():
class SimpleModel(pydantic.BaseModel):
key1_simple: int
key2_simple: float
def func_under_test(x: list[Union[int, SimpleModel]]):
result = []
for item in x:
if isinstance(item, int):
result.append(item + 1)
elif isinstance(item, SimpleModel):
result.append(item.model_copy())
result[-1].key1_simple += 1
result[-1].key2_simple += 1.0
else:
raise ValueError('Unsupported type: %s' % type(item))
return result
expected_response = [2, {'key1_simple': 2, 'key2_simple': 2.0}]
original_args = {'x': [1, {'key1_simple': 1, 'key2_simple': 1.0}]}
actual_response = invoke_function_from_dict_args(
original_args, func_under_test
)
assert isinstance(actual_response, list)
assert len(actual_response) == 2
assert isinstance(actual_response[0], int)
assert isinstance(actual_response[1], SimpleModel)
assert actual_response[0] == expected_response[0]
assert actual_response[1].model_dump() == expected_response[1]
def test_unknown_pydantic_model_argument():
class SimpleModel(pydantic.BaseModel):
key1_simple: int
key2_simple: float
def func_under_test(x: SimpleModel):
return x.model_copy()
original_args = {'x': {'key3_simple': 1, 'key2_simple': 1.0}}
with pytest.raises(errors.UnknownFunctionCallArgumentError):
invoke_function_from_dict_args(original_args, func_under_test)
def test_unknown_pydantic_model_argument_with_union_type():
class SimpleModel1(pydantic.BaseModel):
key1_simple: int
key2_simple: float
class SimpleModel2(pydantic.BaseModel):
key3_simple: str
key4_simple: float
def func_under_test(x: Union[SimpleModel1, SimpleModel2]):
return x.model_copy()
original_args = {'x': {'key5_simple': 1, 'key4_simple': 1.0}}
with pytest.raises(errors.UnknownFunctionCallArgumentError):
invoke_function_from_dict_args(original_args, func_under_test)
def test_unknown_pydantic_model_argument_with_union_type_and_builtin_type():
class SimpleModel1(pydantic.BaseModel):
key1_simple: int
key2_simple: float
def func_under_test(x: Union[SimpleModel1, int]):
return x.model_copy()
original_args = {'x': {'key5_simple': 1, 'key4_simple': 1.0}}
with pytest.raises(errors.UnknownFunctionCallArgumentError):
invoke_function_from_dict_args(original_args, func_under_test)
def test_incompatible_value_and_annotation():
def func_under_test(x: int):
return x + 1
with pytest.raises(errors.UnknownFunctionCallArgumentError):
invoke_function_from_dict_args({'x': {'k': 'v'}}, func_under_test)
with pytest.raises(errors.UnknownFunctionCallArgumentError):
invoke_function_from_dict_args({'x': 'a'}, func_under_test)
with pytest.raises(errors.UnknownFunctionCallArgumentError):
invoke_function_from_dict_args({'x': []}, func_under_test)
with pytest.raises(errors.UnknownFunctionCallArgumentError):
invoke_function_from_dict_args({'x': 1.0}, func_under_test)
with pytest.raises(errors.UnknownFunctionCallArgumentError):
invoke_function_from_dict_args({'x': {}}, func_under_test)
def test_function_invocation_error():
def func_under_test(x: int):
return x / 0
with pytest.raises(errors.FunctionInvocationError):
invoke_function_from_dict_args({'x': 1}, func_under_test)

View File

@@ -0,0 +1,159 @@
# 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.
#
"""Tests for raise_error_for_afc_incompatible_config."""
import pytest
from ... import types
from ..._extra_utils import raise_error_for_afc_incompatible_config
def test_config_is_none():
assert raise_error_for_afc_incompatible_config(None) is None
def test_tool_config_config_unset():
assert (
raise_error_for_afc_incompatible_config(types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=False,
maximum_remote_calls=1,
),
))
is None
)
def test_function_calling_config_unset():
assert (
raise_error_for_afc_incompatible_config(types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=False,
maximum_remote_calls=1,
),
tool_config=types.ToolConfig(),
))
is None
)
def test_compatible_config_afc_disabled():
assert (
raise_error_for_afc_incompatible_config(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=True,
),
tool_config=types.ToolConfig(
function_calling_config=types.FunctionCallingConfig(
stream_function_call_arguments=False,
),
),
)
)
is None
)
def test_compatible_config_stream_function_call_arguments_unset_afc_unset():
assert (
raise_error_for_afc_incompatible_config(
types.GenerateContentConfig(
tool_config=types.ToolConfig(
function_calling_config=types.FunctionCallingConfig(
),
),
)
)
is None
)
def test_compatible_config_stream_function_call_arguments_unset_no_disable_afc():
assert (
raise_error_for_afc_incompatible_config(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
),
tool_config=types.ToolConfig(
function_calling_config=types.FunctionCallingConfig(
),
),
)
)
is None
)
def test_compatible_config_stream_function_call_arguments_unset_disable_afc_true():
assert (
raise_error_for_afc_incompatible_config(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=True,
),
tool_config=types.ToolConfig(
function_calling_config=types.FunctionCallingConfig(
),
),
)
)
is None
)
def test_incompatible_config_stream_function_call_arguments_set_enable_afc():
with pytest.raises(ValueError):
raise_error_for_afc_incompatible_config(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=False,
),
tool_config=types.ToolConfig(
function_calling_config=types.FunctionCallingConfig(
stream_function_call_arguments=True,
),
),
)
)
def test_incompatible_config_stream_function_call_arguments_set_no_afc_config():
with pytest.raises(ValueError):
raise_error_for_afc_incompatible_config(
types.GenerateContentConfig(
tool_config=types.ToolConfig(
function_calling_config=types.FunctionCallingConfig(
stream_function_call_arguments=True,
),
),
)
)
def test_incompatible_config_stream_function_call_arguments_set_no_disable_afc():
with pytest.raises(ValueError):
raise_error_for_afc_incompatible_config(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(),
tool_config=types.ToolConfig(
function_calling_config=types.FunctionCallingConfig(
stream_function_call_arguments=True,
),
),
)
)

View File

@@ -0,0 +1,53 @@
# 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.
#
"""Tests for _extra_utils.should_append_afc_history."""
from ... import types
from ..._extra_utils import should_append_afc_history
def test_should_append_afc_history_with_default_config():
config = types.GenerateContentConfig()
assert should_append_afc_history(config) == True
def test_should_append_afc_history_with_empty_afc_config():
config = types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig()
)
assert should_append_afc_history(config) == True
def test_should_append_afc_history_with_ignore_call_history_true():
config = types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
ignore_call_history=True
)
)
assert should_append_afc_history(config) == False
def test_should_append_afc_history_with_ignore_call_history_false():
config = types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
ignore_call_history=False
)
)
assert should_append_afc_history(config) == True

View File

@@ -0,0 +1,214 @@
# 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.
#
"""Tests for should_disable_afc."""
import pytest
from ... import types
from ..._extra_utils import should_disable_afc
def test_config_is_none():
assert should_disable_afc(None) is False
def test_afc_config_unset():
assert should_disable_afc(types.GenerateContentConfig()) is False
def test_afc_enable_unset_max_0():
assert (
should_disable_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
maximum_remote_calls=0,
),
)
)
is True
)
def test_afc_enable_unset_max_negative():
assert (
should_disable_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
maximum_remote_calls=-1,
),
)
)
is True
)
def test_afc_enable_unset_max_0_0():
assert (
should_disable_afc(
{'automatic_function_calling': {'maximum_remote_calls': 0.0}}
)
is True
)
def test_afc_enable_unset_max_1():
assert (
should_disable_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
maximum_remote_calls=1,
),
)
)
is False
)
def test_afc_enable_unset_max_1_0():
assert (
should_disable_afc(
{'automatic_function_calling': {'maximum_remote_calls': 1.0}}
)
is False
)
def test_afc_enable_false_max_unset():
assert (
should_disable_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=True,
),
)
)
is True
)
def test_afc_enable_false_max_0():
assert (
should_disable_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=True,
maximum_remote_calls=0,
),
)
)
is True
)
def test_afc_enable_false_max_negative():
assert (
should_disable_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=True,
maximum_remote_calls=-1,
),
)
)
is True
)
def test_afc_enable_false_max_0_0():
assert (
should_disable_afc(
{'automatic_function_calling': {'maximum_remote_calls': 0.0}}
)
is True
)
def test_afc_enable_false_max_1():
assert (
should_disable_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=True,
maximum_remote_calls=1,
),
)
)
is True
)
def test_afc_enable_true_max_unset():
assert (
should_disable_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=False,
),
)
)
is False
)
def test_afc_enable_true_max_0():
assert (
should_disable_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=False,
maximum_remote_calls=0,
),
)
)
is True
)
def test_afc_enable_true_max_negative():
assert (
should_disable_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=False,
maximum_remote_calls=-1,
),
)
)
is True
)
def test_afc_enable_true_max_0_0():
assert (
should_disable_afc(
{'automatic_function_calling': {'maximum_remote_calls': 0.0}}
)
is True
)
def test_afc_enable_true_max_1():
assert (
should_disable_afc(
types.GenerateContentConfig(
automatic_function_calling=types.AutomaticFunctionCallingConfig(
disable=False,
maximum_remote_calls=1,
),
)
)
is False
)

View File

@@ -0,0 +1,17 @@
# 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.
#
"""Tests for the Google GenAI SDK's batches module."""

View File

@@ -0,0 +1,77 @@
# 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.
#
"""Tests for batches.cancel()."""
import pytest
from ... import types
from .. import pytest_helper
# Vertex AI batch job name.
_BATCH_JOB_NAME = '6339625664542408704'
_BATCH_JOB_FULL_RESOURCE_NAME = (
'projects/964831358985/locations/us-central1/'
f'batchPredictionJobs/{_BATCH_JOB_NAME}'
)
# MLDev batch operation name.
_MLDEV_BATCH_OPERATION_NAME = 'batches/0yew7plxupyybd7appsrq5vw7w0lp3l79lab'
_INVALID_BATCH_JOB_NAME = 'invalid_name'
# All tests will be run for both Vertex and MLDev.
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_cancel_batch_job',
parameters=types._CancelBatchJobParameters(
name=_BATCH_JOB_NAME,
),
exception_if_mldev='Invalid batch job name',
skip_in_api_mode=('Cannot cancel a batch job multiple times in Vertex'),
),
pytest_helper.TestTableItem(
name='test_cancel_batch_operation',
parameters=types._CancelBatchJobParameters(
name=_MLDEV_BATCH_OPERATION_NAME,
),
exception_if_vertex='Invalid batch job name',
),
pytest_helper.TestTableItem(
name='test_cancel_batch_job_with_invalid_name',
parameters=types._CancelBatchJobParameters(
name=_INVALID_BATCH_JOB_NAME,
),
exception_if_mldev='Invalid batch job name',
exception_if_vertex='Invalid batch job name',
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='batches.cancel',
test_table=test_table,
)
@pytest.mark.asyncio
async def test_async_cancel(client):
if client.vertexai:
name = _BATCH_JOB_NAME
else:
name = _MLDEV_BATCH_OPERATION_NAME
await client.aio.batches.cancel(name=name)

View File

@@ -0,0 +1,123 @@
# 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.
#
"""Tests for batches.create() with invalid source and destinations."""
import pytest
from ... import types
from .. import pytest_helper
_GEMINI_MODEL = 'gemini-2.5-flash'
_EMBEDDING_MODEL = 'gemini-embedding-001'
_DISPLAY_NAME = 'test_batch'
_GENERATE_CONTENT_BQ_OUTPUT_PREFIX = (
'bq://vertex-sdk-dev.unified_genai_tests_batches.generate_content_output'
)
_EMBEDDING_BQ_INPUT_FILE = (
'bq://vertex-sdk-dev.unified_genai_tests_batches.embedding_requests'
)
# All tests will be run for both Vertex and MLDev.
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_union_with_invalid_src',
parameters=types._CreateBatchJobParameters(
model=_GEMINI_MODEL,
src='invalid_src',
config={
'display_name': _DISPLAY_NAME,
'dest': _GENERATE_CONTENT_BQ_OUTPUT_PREFIX,
},
),
exception_if_mldev='Unsupported source',
exception_if_vertex='Unsupported source',
has_union=True,
),
pytest_helper.TestTableItem(
name='test_union_with_invalid_dest',
parameters=types._CreateBatchJobParameters(
model=_EMBEDDING_MODEL,
src=_EMBEDDING_BQ_INPUT_FILE,
config={
'display_name': _DISPLAY_NAME,
'dest': 'invalid_dest',
},
),
exception_if_mldev='not supported in Gemini API',
exception_if_vertex='Unsupported destination',
has_union=True,
),
pytest_helper.TestTableItem(
name='test_create_with_webhook_config',
parameters=types._CreateBatchJobParameters(
model=_GEMINI_MODEL,
src={
'inlined_requests': [
{
'contents': [
{'parts': [{'text': 'say hello'}], 'role': 'user'}
]
},
]
},
config=types.CreateBatchJobConfig(
display_name=_DISPLAY_NAME,
webhook_config=types.WebhookConfig(
uris=['https://example.com/webhook'],
user_metadata={'batch_id': '123'},
),
),
),
exception_if_vertex='not supported in Vertex AI',
),
pytest_helper.TestTableItem(
name='test_create_with_webhook_config_dict',
parameters=types._CreateBatchJobParameters(
model=_GEMINI_MODEL,
src={
'inlined_requests': [
{
'contents': [
{'parts': [{'text': 'say hello'}], 'role': 'user'}
]
},
]
},
config={
'display_name': _DISPLAY_NAME,
'webhook_config': {
'uris': ['https://example.com/webhook'],
},
},
),
exception_if_vertex='not supported in Vertex AI',
),
]
pytestmark = [
pytest.mark.usefixtures('mock_timestamped_unique_name'),
pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='batches.create',
test_table=test_table,
),
]

View File

@@ -0,0 +1,113 @@
# 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.
#
"""Tests for batches.create() with BigQuery source."""
import pytest
from ... import types
from .. import pytest_helper
_GEMINI_MODEL = 'gemini-2.5-flash'
_GEMINI_MODEL_FULL_NAME = 'publishers/google/models/gemini-2.5-flash'
_EMBEDDING_MODEL = 'gemini-embedding-001'
_EMBEDDING_MODEL_FULL_NAME = 'publishers/google/models/gemini-embedding-001'
_DISPLAY_NAME = 'test_batch'
_GENERATE_CONTENT_BQ_INPUT_FILE = (
'bq://vertex-sdk-dev.unified_genai_tests_batches.generate_content_requests'
)
_GENERATE_CONTENT_BQ_OUTPUT_PREFIX = (
'bq://vertex-sdk-dev.unified_genai_tests_batches.generate_content_output'
)
_EMBEDDING_BQ_INPUT_FILE = (
'bq://vertex-sdk-dev.unified_genai_tests_batches.embedding_requests'
)
_EMBEDDING_BQ_OUTPUT_PREFIX = (
'bq://vertex-sdk-dev.unified_genai_tests_batches.embedding_output'
)
# All tests will be run for both Vertex and MLDev.
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_union_generate_content_with_bigquery',
parameters=types._CreateBatchJobParameters(
model=_GEMINI_MODEL_FULL_NAME,
src=_GENERATE_CONTENT_BQ_INPUT_FILE,
),
exception_if_mldev='not supported in Gemini API',
has_union=True,
),
pytest_helper.TestTableItem(
name='test_union_embedding_with_bigquery',
parameters=types._CreateBatchJobParameters(
model='publishers/google/models/text-embedding-004',
src=_EMBEDDING_BQ_INPUT_FILE,
config={
'display_name': _DISPLAY_NAME,
'dest': _EMBEDDING_BQ_OUTPUT_PREFIX,
},
),
exception_if_mldev='not supported in Gemini API',
has_union=True,
),
pytest_helper.TestTableItem(
name='test_generate_content_with_bigquery',
parameters=types._CreateBatchJobParameters(
model=_GEMINI_MODEL_FULL_NAME,
src={
'bigquery_uri': _GENERATE_CONTENT_BQ_INPUT_FILE,
'format': 'bigquery',
},
config={
'display_name': _DISPLAY_NAME,
'dest': {
'bigquery_uri': _GENERATE_CONTENT_BQ_OUTPUT_PREFIX,
'format': 'bigquery',
},
},
),
exception_if_mldev='one of',
),
]
pytestmark = [
pytest.mark.usefixtures('mock_timestamped_unique_name'),
pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='batches.create',
test_table=test_table,
),
]
@pytest.mark.asyncio
async def test_async_create(client):
with pytest_helper.exception_if_mldev(client, ValueError):
batch_job = await client.aio.batches.create(
model=_GEMINI_MODEL,
src=_GENERATE_CONTENT_BQ_INPUT_FILE,
)
assert batch_job.name.startswith('projects/')
assert (
batch_job.model == _GEMINI_MODEL_FULL_NAME
) # Converted to Vertex full name.
assert batch_job.src.bigquery_uri == _GENERATE_CONTENT_BQ_INPUT_FILE
assert batch_job.src.format == 'bigquery'

View File

@@ -0,0 +1,82 @@
# 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.
#
"""Tests for batches.create() with file source."""
import pytest
from ... import types
from .. import pytest_helper
_GEMINI_MODEL = 'gemini-2.5-flash'
_DISPLAY_NAME = 'test_batch'
_MLDEV_GEMINI_MODEL = 'gemini-2.5-flash'
_FILE_NAME = 'files/s0pa54alni6w'
# All tests will be run for both Vertex and MLDev.
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_union_generate_content_with_file',
parameters=types._CreateBatchJobParameters(
model=_MLDEV_GEMINI_MODEL,
src={'file_name': _FILE_NAME},
config={
'display_name': _DISPLAY_NAME,
},
),
exception_if_vertex='not supported',
has_union=True,
),
pytest_helper.TestTableItem(
name='test_generate_content_with_file',
parameters=types._CreateBatchJobParameters(
model=_MLDEV_GEMINI_MODEL,
src={'file_name': _FILE_NAME},
config={
'display_name': _DISPLAY_NAME,
},
),
exception_if_vertex='not supported',
),
]
pytestmark = [
pytest.mark.usefixtures('mock_timestamped_unique_name'),
pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='batches.create',
test_table=test_table,
),
]
@pytest.mark.asyncio
async def test_async_create(client):
with pytest_helper.exception_if_vertex(client, ValueError):
batch_job = await client.aio.batches.create(
model=_MLDEV_GEMINI_MODEL,
src={'file_name': _FILE_NAME},
config={
'display_name': _DISPLAY_NAME,
},
)
assert batch_job.name.startswith('batches/')
assert (
batch_job.model == 'models/' + _MLDEV_GEMINI_MODEL
) # Converted to Gemini full name.

View File

@@ -0,0 +1,125 @@
# 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.
#
"""Tests for batches.create() with Google Cloud Storage source."""
import pytest
from ... import types
from .. import pytest_helper
_GEMINI_MODEL = 'gemini-2.5-flash'
_GEMINI_MODEL_FULL_NAME = 'publishers/google/models/gemini-2.5-flash'
_EMBEDDING_MODEL = 'gemini-embedding-001'
_EMBEDDING_MODEL_FULL_NAME = 'publishers/google/models/gemini-embedding-001'
_DISPLAY_NAME = 'test_batch'
_GENERATE_CONTENT_GCS_INPUT_FILE = (
'gs://unified-genai-tests/batches/input/generate_content_requests.jsonl'
)
_GENERATE_CONTENT_GCS_OUTPUT_PREFIX = 'gs://unified-genai-tests/batches/output'
_EMBEDDING_GCS_INPUT_FILE = (
'gs://unified-genai-tests/batches/input/embedding_requests.jsonl'
)
_EMBEDDING_GCS_OUTPUT_PREFIX = 'gs://unified-genai-tests/batches/output'
# All tests will be run for both Vertex and MLDev.
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_union_generate_content_with_gcs',
parameters=types._CreateBatchJobParameters(
model=_GEMINI_MODEL,
src=_GENERATE_CONTENT_GCS_INPUT_FILE,
),
exception_if_mldev='not supported in Gemini API',
has_union=True,
),
pytest_helper.TestTableItem(
name='test_union_embedding_with_gcs',
parameters=types._CreateBatchJobParameters(
model='text-embedding-004',
src=_EMBEDDING_GCS_INPUT_FILE,
config={
'display_name': _DISPLAY_NAME,
'dest': _EMBEDDING_GCS_OUTPUT_PREFIX,
},
),
exception_if_mldev='not supported in Gemini API',
has_union=True,
),
pytest_helper.TestTableItem(
name='test_union_with_http_options',
parameters=types._CreateBatchJobParameters(
model=_GEMINI_MODEL,
src=_GENERATE_CONTENT_GCS_INPUT_FILE,
config={
'http_options': {
'api_version': 'v1',
'headers': {'test': 'headers'},
},
},
),
exception_if_mldev='not supported in Gemini API',
has_union=True,
),
pytest_helper.TestTableItem(
name='test_generate_content_with_gcs',
parameters=types._CreateBatchJobParameters(
model=_GEMINI_MODEL,
src={
'gcs_uri': [_GENERATE_CONTENT_GCS_INPUT_FILE],
'format': 'jsonl',
},
config={
'display_name': _DISPLAY_NAME,
'dest': {
'gcs_uri': _GENERATE_CONTENT_GCS_OUTPUT_PREFIX,
'format': 'jsonl',
},
},
),
exception_if_mldev='one of',
),
]
pytestmark = [
pytest.mark.usefixtures('mock_timestamped_unique_name'),
pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='batches.create',
test_table=test_table,
),
]
@pytest.mark.asyncio
async def test_async_create(client):
with pytest_helper.exception_if_mldev(client, ValueError):
batch_job = await client.aio.batches.create(
model=_GEMINI_MODEL,
src=_GENERATE_CONTENT_GCS_INPUT_FILE,
)
assert batch_job.name.startswith('projects/')
assert (
batch_job.model == _GEMINI_MODEL_FULL_NAME
) # Converted to Vertex full name.
assert batch_job.src.gcs_uri == [_GENERATE_CONTENT_GCS_INPUT_FILE]
assert batch_job.src.format == 'jsonl'

View File

@@ -0,0 +1,271 @@
# 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.
#
"""Tests for batches.create() with inlined requests."""
import base64
import copy
import datetime
import os
import pytest
from ... import _transformers as t
from ... import types
from .. import pytest_helper
_GEMINI_MODEL = 'gemini-2.5-flash'
_DISPLAY_NAME = 'test_batch'
_MLDEV_GEMINI_MODEL = 'gemini-2.5-flash'
_SAFETY_SETTINGS = [
{
'category': 'HARM_CATEGORY_HATE_SPEECH',
'threshold': 'BLOCK_ONLY_HIGH',
},
{
'category': 'HARM_CATEGORY_DANGEROUS_CONTENT',
'threshold': 'BLOCK_LOW_AND_ABOVE',
},
]
_INLINED_REQUESTS = [
{
'contents': [{
'parts': [{
'text': 'what is the number after 1? return just the number.',
}],
'role': 'user',
}],
'metadata': {
'key': 'request-1',
},
'config': {
'safety_settings': _SAFETY_SETTINGS,
},
},
{
'contents': [{
'parts': [{
'text': 'what is the number after 2? return just the number.',
}],
'role': 'user',
}],
'metadata': {
'key': 'request-2',
},
'config': {
'safety_settings': _SAFETY_SETTINGS,
},
},
]
_INLINED_TEXT_REQUEST_UNION = {
'contents': [{
'parts': [{
'text': 'high',
}],
'role': 'user',
}],
'config': {
'response_modalities': ['TEXT'],
'system_instruction': 'I say high, you say low',
'thinking_config': {
'include_thoughts': True,
'thinking_budget': 4000,
},
},
}
_INLINED_TEXT_REQUEST = copy.deepcopy(_INLINED_TEXT_REQUEST_UNION)
_INLINED_TEXT_REQUEST['config']['system_instruction'] = t.t_content(
'I say high, you say low'
)
_INLINED_TEXT_REQUEST['config']['tools'] = [{'google_search': {}}]
_INLINED_TEXT_REQUEST['config']['tool_config'] = {
'retrieval_config': {'lat_lng': {'latitude': 37.422, 'longitude': -122.084}}
}
_INLINED_IMAGE_REQUEST = {
'contents': [{
'parts': [
{'text': 'What is in this image?'},
{
'file_data': {
'file_uri': (
'https://generativelanguage.googleapis.com/v1beta/files/kje1wewvo85z'
),
'mime_type': 'image/jpeg',
},
},
],
'role': 'user',
}],
'config': {
'temperature': 0.7,
'top_p': 0.9,
'top_k': 10,
},
}
_INLINED_VIDEO_REQUEST = {
'contents': [{
'parts': [
{
'text': 'Summerize this video.',
},
{
'file_data': {
'file_uri': (
'https://generativelanguage.googleapis.com/v1beta/files/tyvaih24jwje'
),
'mime_type': 'video/mp4',
},
'video_metadata': {
'start_offset': '0s',
'end_offset': '5s',
'fps': 3,
},
},
],
'role': 'user',
}],
}
_IMAGE_PNG_FILE_PATH = os.path.abspath(
os.path.join(os.path.dirname(__file__), '../data/google.png')
)
with open(_IMAGE_PNG_FILE_PATH, 'rb') as image_file:
image_bytes = image_file.read()
image_string = base64.b64encode(image_bytes).decode('utf-8')
_INLINED_IMAGE_BLOB_REQUEST = types.InlinedRequest(
contents=[
types.Content(
parts=[
types.Part(text='What is this image about?'),
types.Part(
inline_data=types.Blob(
data=image_string,
mime_type='image/png',
),
),
],
role='user',
)
]
)
# All tests will be run for both Vertex and MLDev.
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_union_with_inlined_request',
parameters=types._CreateBatchJobParameters(
model=_MLDEV_GEMINI_MODEL,
src=_INLINED_REQUESTS,
config={
'display_name': _DISPLAY_NAME,
},
),
exception_if_vertex='not supported in Vertex',
has_union=True,
),
pytest_helper.TestTableItem(
name='test_with_inlined_request',
parameters=types._CreateBatchJobParameters(
model=_MLDEV_GEMINI_MODEL,
src={'inlined_requests': _INLINED_REQUESTS},
config={
'display_name': _DISPLAY_NAME,
},
),
exception_if_vertex='not supported',
),
pytest_helper.TestTableItem(
name='test_with_inlined_request_config',
parameters=types._CreateBatchJobParameters(
model=_MLDEV_GEMINI_MODEL,
src={'inlined_requests': _INLINED_REQUESTS},
config={
'display_name': _DISPLAY_NAME,
},
),
exception_if_vertex='not supported',
),
pytest_helper.TestTableItem(
name='test_union_with_inlined_request_system_instruction',
parameters=types._CreateBatchJobParameters(
model=_MLDEV_GEMINI_MODEL,
src={'inlined_requests': [_INLINED_TEXT_REQUEST_UNION]},
config={
'display_name': _DISPLAY_NAME,
},
),
has_union=True,
exception_if_vertex='not supported',
),
pytest_helper.TestTableItem(
name='test_with_image_file',
parameters=types._CreateBatchJobParameters(
model=_MLDEV_GEMINI_MODEL,
src={'inlined_requests': [_INLINED_IMAGE_REQUEST]},
config={
'display_name': _DISPLAY_NAME,
},
),
exception_if_vertex='not supported',
),
pytest_helper.TestTableItem(
name='test_with_image_blob',
parameters=types._CreateBatchJobParameters(
model=_MLDEV_GEMINI_MODEL,
src={'inlined_requests': [_INLINED_IMAGE_BLOB_REQUEST]},
config={
'display_name': _DISPLAY_NAME,
},
),
exception_if_vertex='not supported',
),
pytest_helper.TestTableItem(
name='test_with_video_file',
parameters=types._CreateBatchJobParameters(
model=_MLDEV_GEMINI_MODEL,
src={'inlined_requests': [_INLINED_VIDEO_REQUEST]},
config={
'display_name': _DISPLAY_NAME,
},
),
exception_if_vertex='not supported',
),
]
pytestmark = [
pytest.mark.usefixtures('mock_timestamped_unique_name'),
pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='batches.create',
test_table=test_table,
),
]
@pytest.mark.asyncio
async def test_async_create(client):
with pytest_helper.exception_if_vertex(client, ValueError):
batch_job = await client.aio.batches.create(
model=_GEMINI_MODEL,
src=_INLINED_REQUESTS,
)
assert batch_job.name.startswith('batches/')
assert (
batch_job.model == 'models/' + _GEMINI_MODEL
) # Converted to Gemini full name.

View File

@@ -0,0 +1,86 @@
# 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.
#
"""Tests for batches.delete()."""
import pytest
from ... import types
from .. import pytest_helper
_BATCH_JOB_NAME = '7085929781874655232'
_BATCH_JOB_FULL_RESOURCE_NAME = (
'projects/964831358985/locations/us-central1/'
f'batchPredictionJobs/{_BATCH_JOB_NAME}'
)
_INVALID_BATCH_JOB_NAME = 'invalid_name'
_MLDEV_BATCH_OPERATION_NAME = 'batches/70h2jo0ic2t1zejyl0p4jgi8mk1gj0wvjusv'
# All tests will be run for both Vertex and MLDev.
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_delete_batch_job',
parameters=types._DeleteBatchJobParameters(
name=_BATCH_JOB_NAME,
),
exception_if_mldev='Invalid batch job name',
skip_in_api_mode=('Cannot cancel a batch job multiple times in Vertex'),
),
pytest_helper.TestTableItem(
name='test_delete_batch_job_full_resource_name',
override_replay_id='test_delete_batch_job',
parameters=types._DeleteBatchJobParameters(
name=_BATCH_JOB_FULL_RESOURCE_NAME,
),
exception_if_mldev='Invalid batch job name',
skip_in_api_mode=('Cannot cancel a batch job multiple times in Vertex'),
),
pytest_helper.TestTableItem(
name='test_delete_batch_job_with_invalid_name',
parameters=types._DeleteBatchJobParameters(
name=_INVALID_BATCH_JOB_NAME,
),
exception_if_mldev='Invalid batch job name',
exception_if_vertex='Invalid batch job name',
),
pytest_helper.TestTableItem(
name='test_delete_batch_operation',
parameters=types._DeleteBatchJobParameters(
name=_MLDEV_BATCH_OPERATION_NAME,
),
exception_if_vertex='Invalid batch job name',
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='batches.delete',
test_table=test_table,
)
@pytest.mark.asyncio
async def test_async_delete(client):
if client.vertexai:
name = _BATCH_JOB_NAME
else:
name = _MLDEV_BATCH_OPERATION_NAME
delete_job = await client.aio.batches.delete(name=name)
assert delete_job

View File

@@ -0,0 +1,157 @@
# 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.
#
"""Tests for batches._create_embeddings()"""
import pytest
from ... import types
from .. import pytest_helper
_MLDEV_EMBEDDING_BATCH_INLINE_OPERATION_NAME = (
'batches/wdx71o8cgbzoa6gg3be1mg7g8ulrhapcjgo3'
)
_MLDEV_EMBEDDING_BATCH_FILE_OPERATION_NAME = (
'batches/507oatd242het8ox60pwsmn7tcmtkrj8itff'
)
_DISPLAY_NAME = 'test_batch'
_MLDEV_EMBEDDING_MODEL = 'gemini-embedding-001'
_EMBED_CONTENT_FILE_NAME = 'files/mq9e3mg3u2y5'
_INLINED_EMBED_CONTENT_REQUESTS = {
'config': {'output_dimensionality': 64},
'contents': [
{
'parts': [{
'text': '1',
}],
},
{
'parts': [{
'text': '2',
}],
},
{
'parts': [{
'text': '3',
}],
},
],
}
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_from_inlined',
parameters=types._CreateEmbeddingsBatchJobParameters(
model=_MLDEV_EMBEDDING_MODEL,
src={'inlined_requests': _INLINED_EMBED_CONTENT_REQUESTS},
config={
'display_name': _DISPLAY_NAME,
},
),
exception_if_vertex='Vertex AI does not support',
),
]
pytestmark = [
pytest.mark.usefixtures('mock_timestamped_unique_name'),
pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='batches.create_embeddings',
test_table=test_table,
http_options={
'api_version': 'v1alpha',
'base_url': (
'https://autopush-generativelanguage.sandbox.googleapis.com'
),
},
),
]
@pytest.mark.asyncio
async def test_async_from_inline(client):
with pytest_helper.exception_if_vertex(client, ValueError):
batch_job = await client.aio.batches.create_embeddings(
model=_MLDEV_EMBEDDING_MODEL,
src={'inlined_requests': _INLINED_EMBED_CONTENT_REQUESTS},
)
assert batch_job.name.startswith('batches/')
def test_from_file(client):
"""Tests creating a batch job with an embedding file name."""
with pytest_helper.exception_if_vertex(client, ValueError):
batch_job = client.batches.create_embeddings(
model=_MLDEV_EMBEDDING_MODEL,
src={'file_name': _EMBED_CONTENT_FILE_NAME},
config={
'display_name': _DISPLAY_NAME,
},
)
assert batch_job.name.startswith('batches/')
assert batch_job.model == 'models/' + _MLDEV_EMBEDDING_MODEL
@pytest.mark.asyncio
async def test_async_from_file(client):
with pytest_helper.exception_if_vertex(client, ValueError):
batch_job = await client.aio.batches.create_embeddings(
model=_MLDEV_EMBEDDING_MODEL,
src={'file_name': _EMBED_CONTENT_FILE_NAME},
config={
'display_name': _DISPLAY_NAME,
},
)
assert batch_job.name.startswith('batches/')
assert (
batch_job.model == 'models/' + _MLDEV_EMBEDDING_MODEL
) # Converted to Gemini full name.
def test_get_inline(client):
"""Tests getting a batch job that used inline requests."""
with pytest_helper.exception_if_vertex(client, ValueError):
name = _MLDEV_EMBEDDING_BATCH_INLINE_OPERATION_NAME
batch_job = client.batches.get(name=name)
assert batch_job.dest.inlined_embed_content_responses is not None
@pytest.mark.asyncio
async def test_async_get_inline(client):
with pytest_helper.exception_if_vertex(client, ValueError):
name = _MLDEV_EMBEDDING_BATCH_INLINE_OPERATION_NAME
batch_job = await client.aio.batches.get(name=name)
assert batch_job.dest.inlined_embed_content_responses is not None
def test_get_file(client):
"""Tests getting a batch job that used a file source."""
with pytest_helper.exception_if_vertex(client, ValueError):
name = _MLDEV_EMBEDDING_BATCH_FILE_OPERATION_NAME
batch_job = client.batches.get(name=name)
assert batch_job.dest.file_name is not None
@pytest.mark.asyncio
async def test_async_get_file(client):
with pytest_helper.exception_if_vertex(client, ValueError):
name = _MLDEV_EMBEDDING_BATCH_FILE_OPERATION_NAME
batch_job = await client.aio.batches.get(name=name)
assert batch_job.dest.file_name is not None

View File

@@ -0,0 +1,78 @@
# 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.
#
"""Tests for batches.get()."""
import pytest
from ... import types
from .. import pytest_helper
# Vertex AI batch job name.
_BATCH_JOB_NAME = '5798522612028014592'
_BATCH_JOB_FULL_RESOURCE_NAME = (
'projects/964831358985/locations/us-central1/'
f'batchPredictionJobs/{_BATCH_JOB_NAME}'
)
# MLDev batch operation name.
_MLDEV_BATCH_OPERATION_NAME = 'batches/z2p8ksus4lyxt25rntl3fpd67p2niw4hfij5'
_INVALID_BATCH_JOB_NAME = 'invalid_name'
# All tests will be run for both Vertex and MLDev.
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_get_batch_job',
parameters=types._GetBatchJobParameters(
name=_BATCH_JOB_NAME,
),
exception_if_mldev='Invalid batch job name',
),
pytest_helper.TestTableItem(
name='test_get_batch_operation',
parameters=types._GetBatchJobParameters(
name=_MLDEV_BATCH_OPERATION_NAME,
),
exception_if_vertex='Invalid batch job name',
),
pytest_helper.TestTableItem(
name='test_get_batch_job_with_invalid_name',
parameters=types._GetBatchJobParameters(
name=_INVALID_BATCH_JOB_NAME,
),
exception_if_mldev='Invalid batch job name',
exception_if_vertex='Invalid batch job name',
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='batches.get',
test_table=test_table,
)
@pytest.mark.asyncio
async def test_async_get(client):
if client.vertexai:
name = _BATCH_JOB_NAME
else:
name = _MLDEV_BATCH_OPERATION_NAME
batch_job = await client.aio.batches.get(name=name)
assert batch_job

View File

@@ -0,0 +1,79 @@
# 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.
#
"""Tests for batches.list()."""
import pytest
from ... import errors
from ... import types
from .. import pytest_helper
# All tests will be run for both Vertex and MLDev.
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_list_batch_jobs',
parameters=types._ListBatchJobsParameters(
config=types.ListBatchJobsConfig(
page_size=5,
),
),
),
pytest_helper.TestTableItem(
name='test_list_batch_jobs_with_config',
parameters=types._ListBatchJobsParameters(
config=types.ListBatchJobsConfig(
filter='display_name:"genai_*"',
page_size=5,
),
),
exception_if_mldev='filter parameter is not supported in Gemini API',
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='batches.list',
test_table=test_table,
)
def test_pager(client):
batch_jobs = client.batches.list(config={'page_size': 10})
assert 'content-type' in batch_jobs.sdk_http_response.headers
assert batch_jobs.name == 'batch_jobs'
assert batch_jobs.page_size == 10
assert len(batch_jobs) <= 10
# Iterate through the first page. Otherwise, too many batches are returned.
for _ in batch_jobs.page:
pass
@pytest.mark.asyncio
async def test_async_pager(client):
batch_jobs = await client.aio.batches.list(config={'page_size': 10})
assert 'Content-Type' in batch_jobs.sdk_http_response.headers
assert batch_jobs.name == 'batch_jobs'
assert batch_jobs.page_size == 10
assert len(batch_jobs) <= 10
# Iterate through the first page. Otherwise, too many batches are returned.
for _ in batch_jobs.page:
pass

View File

@@ -0,0 +1,17 @@
# 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.
#
"""Tests for the Google GenAI SDK's caches module."""

View File

@@ -0,0 +1,29 @@
# 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.
#
"""Constants for caches tests."""
CACHED_CONTENT_NAME_MLDEV = 'cachedContents/o239k1gxzz0juy9wqstndhncr85krehehf551hqh'
CACHED_CONTENT_NAME_VERTEX = 'cachedContents/2164089915711684608'
VERTEX_HTTP_OPTIONS = {
'api_version': 'v1beta1',
'base_url': 'https://us-central1-aiplatform.googleapis.com/',
}
MLDEV_HTTP_OPTIONS = {
'api_version': 'v1beta',
'base_url': 'https://generativelanguage.googleapis.com/',
}

View File

@@ -0,0 +1,210 @@
# 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.
#
from copy import deepcopy
import datetime
import pytest
import sys
from ... import types
from .. import pytest_helper
from ... import _transformers as t
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI = types._CreateCachedContentParameters(
model='gemini-2.5-flash',
config={
'contents': [
types.Content(
role='user',
parts=[
types.Part(
fileData=types.FileData(
fileUri='gs://cloud-samples-data/generative-ai/pdf/2312.11805v3.pdf',
mimeType='application/pdf',
)
),
types.Part(
fileData=types.FileData(
fileUri='gs://cloud-samples-data/generative-ai/pdf/2403.05530.pdf',
mimeType='application/pdf',
)
),
],
)
],
'system_instruction': t.t_content('What is the sum of the two pdfs?'),
'display_name': 'test cache',
'ttl': '86400s',
},
)
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE = types._CreateCachedContentParameters(
model='gemini-2.5-flash',
config={
'contents': [
types.Content(
role='user',
parts=[
types.Part(
fileData=types.FileData(
mimeType='application/pdf',
fileUri='https://generativelanguage.googleapis.com/v1beta/files/v200dhvn15h7',
)
)
],
)
],
'system_instruction': t.t_content('What is the sum of the two pdfs?'),
'display_name': 'test cache',
'ttl': '86400s',
},
)
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_PARTIAL_MODEL_1 = deepcopy(
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI
)
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_PARTIAL_MODEL_1.model = (
'models/gemini-2.5-flash'
)
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_PARTIAL_MODEL_2 = deepcopy(
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI
)
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_PARTIAL_MODEL_2.model = (
'publishers/google/models/gemini-2.5-flash'
)
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE_PARTIAL_MODEL_1 = deepcopy(
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE
)
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE_PARTIAL_MODEL_1.model = (
'models/gemini-2.5-flash'
)
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_CMEK = deepcopy(
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI
)
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_CMEK.config.kms_key_name = (
'projects/test-project/locations/us-central1/keyRings/test-keyring/cryptoKeys/test-key'
)
if sys.version_info >= (3, 11):
_EXPIRE_TIME = datetime.datetime.fromisoformat('2025-12-20T00:00:00Z')
else:
_EXPIRE_TIME = datetime.datetime.fromisoformat('2025-12-20T00:00:00+00:00')
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_EXPIRE_TIME = deepcopy(
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI
)
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_EXPIRE_TIME.config.ttl = None
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_EXPIRE_TIME.config.expire_time = (
_EXPIRE_TIME
)
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_EXPIRE_TIME.config.display_name = (
'test cache'
)
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE_EXPIRE_TIME = deepcopy(
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE
)
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE_EXPIRE_TIME.config.ttl = None
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE_EXPIRE_TIME.config.expire_time = (
_EXPIRE_TIME
)
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE_EXPIRE_TIME.config.display_name = (
'test cache'
)
# Replay mode is not supported for caches tests due to the error message
# inconsistency in api and replay mode.
# To run api mode tests, use the following steps:
# 1. First create the resource.
# sh run_tests.sh pytest -s tests/caches/test_create.py --mode=api
# 1.1 If mldev test_create fails, update the uploaded file using this colab
# https://colab.sandbox.google.com/drive/1Fv6KGSs0cg6tlpcUHdsclHussXMEGOXk#scrollTo=RSKmFPx00MVL.
# 2. Find the resource name in debugging print and change the resource name constants.py.
# 3. Run and record get and update tests.
# sh run_tests.sh pytest -s tests/caches/test_get.py --mode=api && sh run_tests.sh pytest -s tests/caches/test_update.py --mode=api && sh run_tests.sh pytest -s tests/caches/test_delete.py --mode=api
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_caches_create_with_gcs_uri',
exception_if_mldev='INVALID_ARGUMENT',
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI,
),
pytest_helper.TestTableItem(
name='test_caches_create_with_gcs_uri_cmek',
exception_if_mldev='not supported',
exception_if_vertex='INVALID_ARGUMENT', # The key is invalid.
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_CMEK,
),
pytest_helper.TestTableItem(
name='test_caches_create_with_gcs_uri_expire_time',
exception_if_mldev='INVALID_ARGUMENT',
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_EXPIRE_TIME,
),
pytest_helper.TestTableItem(
name='test_caches_create_model_partial_path_1',
exception_if_mldev='INVALID_ARGUMENT',
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_PARTIAL_MODEL_1,
),
pytest_helper.TestTableItem(
name='test_caches_create_model_partial_path_2',
exception_if_mldev='404',
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_PARTIAL_MODEL_2,
),
pytest_helper.TestTableItem(
name='test_caches_create_with_googleai_file',
exception_if_vertex='Internal',
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE,
skip_in_api_mode='Create is not reproducible in the API mode.',
),
pytest_helper.TestTableItem(
name='test_caches_create_with_googleai_file_expire_time',
exception_if_vertex='Internal',
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE_EXPIRE_TIME,
skip_in_api_mode='Create is not reproducible in the API mode.',
),
pytest_helper.TestTableItem(
name='test_caches_create_with_googleai_file_model_partial_path_1',
exception_if_vertex='Internal',
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE_PARTIAL_MODEL_1,
skip_in_api_mode='Create is not reproducible in the API mode.',
),
]
pytestmark = [
pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='caches.create',
test_table=test_table,
),
]
pytest_plugins = ('pytest_asyncio',)
@pytest.mark.asyncio
async def test_async_googleai_file_create(client):
if client._api_client.vertexai:
with pytest.raises(Exception):
await client.aio.caches.create(
model=_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE.model,
config=_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE.config,
)
else:
await client.aio.caches.create(
model=_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE.model,
config=_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE.config,
)

View File

@@ -0,0 +1,105 @@
# 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 copy
import pytest
from ... import types
from .. import pytest_helper
from . import constants
from ... import _transformers as t
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI = types._CreateCachedContentParameters(
model='gemini-2.5-flash',
config={
'contents': [
types.Content(
role='user',
parts=[
types.Part(
fileData=types.FileData(
fileUri='gs://cloud-samples-data/generative-ai/pdf/2312.11805v3.pdf',
mimeType='application/pdf',
)
),
types.Part(
fileData=types.FileData(
fileUri='gs://cloud-samples-data/generative-ai/pdf/2403.05530.pdf',
mimeType='application/pdf',
)
),
],
)
],
'system_instruction': t.t_content('What is the sum of the two pdfs?'),
'display_name': 'test cache',
'ttl': '86400s',
'http_options': constants.VERTEX_HTTP_OPTIONS,
},
)
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE = types._CreateCachedContentParameters(
model='gemini-2.5-flash',
config={
'contents': [
types.Content(
role='user',
parts=[
types.Part(
fileData=types.FileData(
mimeType='application/pdf',
fileUri='https://generativelanguage.googleapis.com/v1beta/files/v200dhvn15h7',
)
)
],
)
],
'system_instruction': t.t_content('What is the sum of the two pdfs?'),
'display_name': 'test cache',
'ttl': '86400s',
'http_options': constants.MLDEV_HTTP_OPTIONS,
},
)
# Replay mode is not supported for caches tests due to the error message
# inconsistency in api and replay mode.
# To run api mode tests, use the following steps:
# 1. First create the resource.
# sh run_tests.sh pytest -s tests/caches/test_create.py --mode=api
# 1.1 If mldev test_create fails, update the uploaded file using this colab
# https://colab.sandbox.google.com/drive/1Fv6KGSs0cg6tlpcUHdsclHussXMEGOXk#scrollTo=RSKmFPx00MVL.
# 2. Find the resource name in debugging print and change the resource name constants.py.
# 3. Run and record get and update tests.
# sh run_tests.sh pytest -s tests/caches/test_get.py --mode=api && sh run_tests.sh pytest -s tests/caches/test_update.py --mode=api && sh run_tests.sh pytest -s tests/caches/test_delete.py --mode=api
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_caches_create_with_gcs_uri',
exception_if_mldev='404',
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI,
),
pytest_helper.TestTableItem(
name='test_caches_create_with_googleai_file',
exception_if_vertex='404',
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE,
skip_in_api_mode='Create is not reproducible in the API mode.',
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='caches.create',
test_table=test_table,
)
pytest_plugins = ('pytest_asyncio',)

View File

@@ -0,0 +1,54 @@
# 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 pytest
from ... import types
from .. import pytest_helper
from . import constants
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
skip_in_api_mode='Delete is not reproducible in the API mode.',
name='test_caches_delete_with_vertex_cache_name',
exception_if_mldev='PERMISSION_DENIED',
parameters=types._DeleteCachedContentParameters(
name=constants.CACHED_CONTENT_NAME_VERTEX,
),
),
pytest_helper.TestTableItem(
skip_in_api_mode='Delete is not reproducible in the API mode.',
name='test_caches_delete_with_mldev_cache_name',
exception_if_vertex='NOT_FOUND',
parameters=types._DeleteCachedContentParameters(
name=constants.CACHED_CONTENT_NAME_MLDEV,
),
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='caches.delete',
test_table=test_table,
)
@pytest.mark.asyncio
async def test_async_delete(client):
if client._api_client.vertexai:
await client.aio.caches.delete(name=constants.CACHED_CONTENT_NAME_VERTEX)
else:
await client.aio.caches.delete(name=constants.CACHED_CONTENT_NAME_MLDEV)

View File

@@ -0,0 +1,52 @@
# 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 pytest
from ... import types
from .. import pytest_helper
from . import constants
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
skip_in_api_mode='Delete is not reproducible in the API mode.',
name='test_caches_delete_with_vertex_cache_name',
exception_if_mldev='404',
parameters=types._DeleteCachedContentParameters(
name=constants.CACHED_CONTENT_NAME_VERTEX,
config={
'http_options': constants.VERTEX_HTTP_OPTIONS,
},
),
),
pytest_helper.TestTableItem(
skip_in_api_mode='Delete is not reproducible in the API mode.',
name='test_caches_delete_with_mldev_cache_name',
exception_if_vertex='404',
parameters=types._DeleteCachedContentParameters(
name=constants.CACHED_CONTENT_NAME_MLDEV,
config={
'http_options': constants.MLDEV_HTTP_OPTIONS,
},
),
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='caches.delete',
test_table=test_table,
)

View File

@@ -0,0 +1,94 @@
# 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 pytest
from ... import types
from .. import pytest_helper
from . import constants
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
skip_in_api_mode='Get is not reproducible in the API mode.',
name='test_caches_get_with_vertex_cache_name',
exception_if_mldev='PERMISSION_DENIED',
parameters=types._GetCachedContentParameters(
name=constants.CACHED_CONTENT_NAME_VERTEX,
),
),
pytest_helper.TestTableItem(
skip_in_api_mode='Get is not reproducible in the API mode.',
name='test_caches_get_with_mldev_cache_name',
exception_if_vertex='NOT_FOUND',
parameters=types._GetCachedContentParameters(
name=constants.CACHED_CONTENT_NAME_MLDEV,
),
),
pytest_helper.TestTableItem(
skip_in_api_mode='Get is not reproducible in the API mode.',
name='test_caches_get_with_mldev_cache_name_partial_1',
exception_if_vertex='NOT_FOUND',
parameters=types._GetCachedContentParameters(
name='cachedContents/o239k1gxzz0juy9wqstndhncr85krehehf551hqh'
),
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='caches.get',
test_table=test_table,
)
@pytest.mark.asyncio
async def test_async_get(client):
if client._api_client.vertexai:
response = await client.aio.caches.get(
name=constants.CACHED_CONTENT_NAME_VERTEX
)
assert response
else:
await client.aio.caches.get(name=constants.CACHED_CONTENT_NAME_MLDEV)
def test_different_cache_name_formats(client):
if client._api_client.vertexai:
response1 = client.caches.get(
name='projects/964831358985/locations/us-central1/cachedContents/2164089915711684608'
)
assert response1
response2 = client.caches.get(
name='locations/us-central1/cachedContents/2164089915711684608'
)
assert response2
response3 = client.caches.get(
name='cachedContents/2164089915711684608'
)
assert response3
response4 = client.caches.get(
name='2164089915711684608'
)
assert response4
else:
response1 = client.caches.get(
name='cachedContents/o239k1gxzz0juy9wqstndhncr85krehehf551hqh'
)
assert response1
response2 = client.caches.get(
name='o239k1gxzz0juy9wqstndhncr85krehehf551hqh'
)
assert response2

View File

@@ -0,0 +1,52 @@
# 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 pytest
from ... import types
from .. import pytest_helper
from . import constants
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
skip_in_api_mode='Get is not reproducible in the API mode.',
name='test_caches_get_with_vertex_cache_name',
exception_if_mldev='404',
parameters=types._GetCachedContentParameters(
name=constants.CACHED_CONTENT_NAME_VERTEX,
config={
'http_options': constants.VERTEX_HTTP_OPTIONS,
},
),
),
pytest_helper.TestTableItem(
skip_in_api_mode='Get is not reproducible in the API mode.',
name='test_caches_get_with_mldev_cache_name',
exception_if_vertex='404',
parameters=types._GetCachedContentParameters(
name=constants.CACHED_CONTENT_NAME_MLDEV,
config={
'http_options': constants.MLDEV_HTTP_OPTIONS,
},
),
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='caches.get',
test_table=test_table,
)

View File

@@ -0,0 +1,68 @@
# 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.
#
"""Test caches list method."""
import pytest
from ... import types
from .. import pytest_helper
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
skip_in_api_mode='List is not reproducible in the API mode.',
name='test_caches_list',
parameters=types._ListCachedContentsParameters(config={'page_size': 3}),
),
]
pytestmark = [
pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='caches.list',
test_table=test_table,
)
]
def test_pager(client):
cached_contents = client.caches.list(config={'page_size': 2})
assert 'content-type' in cached_contents.sdk_http_response.headers
assert cached_contents.name == 'cached_contents'
assert cached_contents.page_size == 2
assert len(cached_contents) <= 2
# Iterate through all the pages. Then next_page() should raise an exception.
for _ in cached_contents:
pass
with pytest.raises(IndexError, match='No more pages to fetch.'):
cached_contents.next_page()
@pytest.mark.asyncio
async def test_async_pager(client):
cached_contents = await client.aio.caches.list(config={'page_size': 2})
assert 'Content-Type' in cached_contents.sdk_http_response.headers
assert cached_contents.name == 'cached_contents'
assert cached_contents.page_size == 2
assert len(cached_contents) <= 2
# Iterate through all the pages. Then next_page() should raise an exception.
async for _ in cached_contents:
pass
with pytest.raises(IndexError, match='No more pages to fetch.'):
await cached_contents.next_page()

View File

@@ -0,0 +1,70 @@
# 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 pytest
from ... import types
from .. import pytest_helper
from . import constants
_VERTEX_UPDATE_PARAMETERS = types._UpdateCachedContentParameters(
name=constants.CACHED_CONTENT_NAME_VERTEX,
config={
'ttl': '7600s',
},
)
_MLDEV_UPDATE_PARAMETERS = types._UpdateCachedContentParameters(
name=constants.CACHED_CONTENT_NAME_MLDEV,
config={
'ttl': '7600s',
},
)
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
skip_in_api_mode='Update has permission issues in the API mode.',
name='test_caches_update_with_vertex_cache_name',
exception_if_mldev='PERMISSION_DENIED',
parameters=_VERTEX_UPDATE_PARAMETERS,
),
pytest_helper.TestTableItem(
skip_in_api_mode='Update has permission issues in the API mode.',
name='test_caches_update_with_mldev_cache_name',
exception_if_vertex='NOT_FOUND',
parameters=_MLDEV_UPDATE_PARAMETERS,
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='caches.update',
test_table=test_table,
)
@pytest.mark.asyncio
async def test_async_update(client):
if client._api_client.vertexai:
response = await client.aio.caches.update(
name=_VERTEX_UPDATE_PARAMETERS.name,
config=_VERTEX_UPDATE_PARAMETERS.config,
)
assert response
else:
await client.aio.caches.update(
name=_MLDEV_UPDATE_PARAMETERS.name,
config=_MLDEV_UPDATE_PARAMETERS.config,
)

View File

@@ -0,0 +1,58 @@
# 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 pytest
from ... import types
from .. import pytest_helper
from . import constants
_VERTEX_UPDATE_PARAMETERS = types._UpdateCachedContentParameters(
name=constants.CACHED_CONTENT_NAME_VERTEX,
config={
'ttl': '7600s',
'http_options': constants.VERTEX_HTTP_OPTIONS,
},
)
_MLDEV_UPDATE_PARAMETERS = types._UpdateCachedContentParameters(
name=constants.CACHED_CONTENT_NAME_MLDEV,
config={
'ttl': '7600s',
'http_options': constants.MLDEV_HTTP_OPTIONS,
},
)
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
skip_in_api_mode='Update has permission issues in the API mode.',
name='test_caches_update_with_vertex_cache_name',
exception_if_mldev='404',
parameters=_VERTEX_UPDATE_PARAMETERS,
),
pytest_helper.TestTableItem(
skip_in_api_mode='Update has permission issues in the API mode.',
name='test_caches_update_with_mldev_cache_name',
exception_if_vertex='404',
parameters=_MLDEV_UPDATE_PARAMETERS,
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='caches.update',
test_table=test_table,
)

View File

@@ -0,0 +1 @@
"""Tests for the Google GenAI SDK's chats module."""

View File

@@ -0,0 +1,598 @@
# 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.
#
from unittest import mock
import pytest
from ... import chats
from ... import client
from ... import models
from ... import types
AFC_HISTORY = [
types.Content(
role='user',
parts=[types.Part.from_text(text='afc input')],
),
types.Content(
role='model',
parts=[
types.Part(
function_call=types.FunctionCall(
name='foo', args={'bar': 'baz'}
)
)
],
),
]
pytest_plugins = 'pytest_asyncio'
@pytest.fixture
def mock_api_client(vertexai=False):
api_client = mock.MagicMock(spec=client.ApiClient)
api_client.api_key = 'TEST_API_KEY'
api_client._host = lambda: 'test_host'
api_client._http_options = {'headers': {}} # Ensure headers exist
api_client.vertexai = vertexai
return api_client
@pytest.fixture
def mock_generate_content_with_empty_text_part():
with mock.patch.object(
models.Models, 'generate_content'
) as mock_generate_content:
mock_generate_content.return_value = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
role='model',
parts=[types.Part(text='')],
)
)
]
)
yield mock_generate_content
@pytest.fixture
def mock_generate_content_empty_content():
with mock.patch.object(
models.Models, 'generate_content'
) as mock_generate_content:
mock_generate_content.return_value = types.GenerateContentResponse(
candidates=[]
)
yield mock_generate_content
@pytest.fixture
def mock_generate_content_stream_with_empty_text_part():
with mock.patch.object(
models.Models, 'generate_content_stream'
) as mock_generate_content:
mock_generate_content.return_value = [
types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
role='model',
parts=[types.Part(text='')],
),
finish_reason=types.FinishReason.STOP,
)
]
)
]
yield mock_generate_content
@pytest.fixture
def mock_generate_content_stream_empty_content():
with mock.patch.object(
models.Models, 'generate_content_stream'
) as mock_generate_content:
mock_generate_content.return_value = [
types.GenerateContentResponse(candidates=[])
]
yield mock_generate_content
@pytest.fixture
def mock_generate_content_afc_history():
with mock.patch.object(
models.Models, 'generate_content'
) as mock_generate_content:
mock_generate_content.return_value = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
role='model',
parts=[types.Part.from_text(text='afc output')],
)
)
],
automatic_function_calling_history=AFC_HISTORY,
)
yield mock_generate_content
@pytest.fixture
def mock_generate_content_stream_afc_history():
with mock.patch.object(
models.Models, 'generate_content_stream'
) as mock_generate_content:
mock_generate_content.return_value = [
types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
role='model',
parts=[types.Part.from_text(text='afc output')],
),
finish_reason=types.FinishReason.STOP,
)
],
automatic_function_calling_history=AFC_HISTORY,
)
]
yield mock_generate_content
def test_history_start_with_valid_model_content():
history = [
types.Content(
role='model',
parts=[types.Part.from_text(text='Hello there! how can I help you?')],
),
types.Content(role='user', parts=[types.Part.from_text(text='Hello')]),
]
models_module = models.Models(mock_api_client)
chats_module = chats.Chats(modules=models_module)
chat = chats_module.create(model='gemini-2.5-flash', history=history)
assert chat.get_history() == history
assert chat.get_history(curated=True) == history
def test_history_start_with_invalid_model_content():
history = [
types.Content(
role='model',
parts=[],
),
types.Content(role='user', parts=[types.Part.from_text(text='Hello')]),
]
models_module = models.Models(mock_api_client)
chats_module = chats.Chats(modules=models_module)
chat = chats_module.create(model='gemini-2.5-flash', history=history)
assert chat.get_history() == history
assert chat.get_history(curated=True) == [types.Content(role='user', parts=[types.Part.from_text(text='Hello')])]
def test_history_with_consecutive_valid_user_inputs():
history = [
types.Content(
role='user',
parts=[types.Part.from_text(text='user input 1')],
),
types.Content(
role='user',
parts=[types.Part.from_text(text='user input 2')],
),
]
models_module = models.Models(mock_api_client)
chats_module = chats.Chats(modules=models_module)
chat = chats_module.create(model='gemini-2.5-flash', history=history)
assert chat.get_history() == history
assert chat.get_history(curated=True) == history
def test_history_with_valid_and_invalid_user_inputs():
history = [
types.Content(
role='user',
parts=[types.Part.from_text(text='user input 1')],
),
types.Content(
role='user',
parts=[], # invalid content
),
types.Content(
role='user',
parts=[types.Part.from_text(text='user input 2')],
),
]
models_module = models.Models(mock_api_client)
chats_module = chats.Chats(modules=models_module)
chat = chats_module.create(model='gemini-2.5-flash', history=history)
assert chat.get_history() == history
assert chat.get_history(curated=True) == history
def test_history_with_consecutive_valid_model_outputs():
history = [
types.Content(
role='model',
parts=[types.Part.from_text(text='model output 1')],
),
types.Content(
role='model',
parts=[types.Part.from_text(text='model output 2')],
),
]
models_module = models.Models(mock_api_client)
chats_module = chats.Chats(modules=models_module)
chat = chats_module.create(model='gemini-2.5-flash', history=history)
assert chat.get_history() == history
assert chat.get_history(curated=True) == history
def test_history_with_valid_and_invalid_model_output():
history = [
types.Content(
role='model',
parts=[types.Part.from_text(text='model output 1')],
),
types.Content(
role='model',
parts=[], # invalid content
),
types.Content(
role='model',
parts=[types.Part.from_text(text='model output 2')],
),
]
models_module = models.Models(mock_api_client)
chats_module = chats.Chats(modules=models_module)
chat = chats_module.create(model='gemini-2.5-flash', history=history)
assert chat.get_history() == history
assert chat.get_history(curated=True) == []
def test_history_end_with_user_input():
history = [
types.Content(
role='user',
parts=[types.Part.from_text(text='user input 1')],
),
types.Content(
role='model',
parts=[types.Part.from_text(text='model output')],
),
types.Content(
role='user',
parts=[types.Part.from_text(text='user input 2')],
),
]
models_module = models.Models(mock_api_client)
chats_module = chats.Chats(modules=models_module)
chat = chats_module.create(model='gemini-2.5-flash', history=history)
assert chat.get_history() == history
assert chat.get_history(curated=True) == history
def test_unrecognized_role_in_history():
history = [
types.Content(role='user', parts=[types.Part.from_text(text='Hello')]),
types.Content(
role='invalid_role',
parts=[types.Part.from_text(text='Hello there! how can I help you?')],
),
]
models_module = models.Models(mock_api_client)
chats_module = chats.Chats(modules=models_module)
with pytest.raises(ValueError) as e:
chat = chats_module.create(model='gemini-2.5-flash', history=history)
assert 'Role must be user or model' in str(e)
def test_sync_chat_create():
history = [
types.Content(
role='user', parts=[types.Part.from_text(text='user input turn 1')]
),
types.Content(
role='model',
parts=[types.Part.from_text(text='model output turn 1')],
),
types.Content(
role='model',
parts=[types.Part.from_text(text='model output turn 1')],
),
types.Content(
role='model',
parts=[types.Part.from_text(text='user input turn 2')],
),
]
models_module = models.Models(mock_api_client)
chats_module = chats.Chats(modules=models_module)
chat = chats_module.create(model='gemini-2.5-flash', history=history)
assert chat.get_history() == history
assert chat.get_history(curated=True) == history
def test_async_chat_create():
history = [
types.Content(
role='user', parts=[types.Part.from_text(text='user input turn 1')]
),
types.Content(
role='model',
parts=[types.Part.from_text(text='model output turn 1')],
),
types.Content(
role='model',
parts=[types.Part.from_text(text='model output turn 1')],
),
types.Content(
role='model',
parts=[types.Part.from_text(text='user input turn 2')],
),
types.Content(
role='model',
parts=[types.Part.from_text(text='model output turn 2')],
),
]
models_module = models.AsyncModels(mock_api_client)
chats_module = chats.AsyncChats(modules=models_module)
chat = chats_module.create(model='gemini-2.5-flash', history=history)
assert chat.get_history() == history
assert chat.get_history(curated=True) == history
def test_sync_chat_create_with_history_dict():
history = [
{'role': 'user', 'parts': [{'text': 'user input turn 1'}]},
{'role': 'model', 'parts': [{'text': 'model output turn 1'}]},
{'role': 'user', 'parts': [{'text': 'user input turn 2'}]},
{'role': 'model', 'parts': [{'text': 'model output turn 2'}]},
]
models_module = models.Models(mock_api_client)
chats_module = chats.Chats(modules=models_module)
chat = chats_module.create(model='gemini-2.5-flash', history=history)
expected_history = [
types.Content(
role='user', parts=[types.Part.from_text(text='user input turn 1')]
),
types.Content(
role='model',
parts=[types.Part.from_text(text='model output turn 1')],
),
types.Content(
role='user', parts=[types.Part.from_text(text='user input turn 2')]
),
types.Content(
role='model',
parts=[types.Part.from_text(text='model output turn 2')],
),
]
assert chat.get_history() == expected_history
assert chat.get_history(curated=True) == expected_history
def test_async_chat_create_with_history_dict():
history = [
{'role': 'user', 'parts': [{'text': 'user input turn 1'}]},
{'role': 'model', 'parts': [{'text': 'model output turn 1'}]},
{'role': 'user', 'parts': [{'text': 'user input turn 2'}]},
{'role': 'model', 'parts': [{'text': 'model output turn 2'}]},
]
models_module = models.AsyncModels(mock_api_client)
chats_module = chats.AsyncChats(modules=models_module)
chat = chats_module.create(model='gemini-2.5-flash', history=history)
expected_history = [
types.Content(
role='user', parts=[types.Part.from_text(text='user input turn 1')]
),
types.Content(
role='model',
parts=[types.Part.from_text(text='model output turn 1')],
),
types.Content(
role='user', parts=[types.Part.from_text(text='user input turn 2')]
),
types.Content(
role='model',
parts=[types.Part.from_text(text='model output turn 2')],
),
]
assert chat.get_history() == expected_history
assert chat.get_history(curated=True) == expected_history
def test_history_with_invalid_turns():
valid_input = types.Content(
role='user', parts=[types.Part.from_text(text='Hello')]
)
valid_output = [
types.Content(
role='model',
parts=[types.Part.from_text(text='Hello there! how can I help you?')],
),
types.Content(
role='model',
parts=[types.Part.from_text(text='Hello there! how can I help you?')],
),
]
invalid_input = types.Content(
role='user',
parts=[
types.Part.from_text(text='a input will be rejected by the model')
],
)
invalid_output = types.Content(
role='model',
parts=[],
)
comprehensive_history = []
comprehensive_history.append(valid_input)
comprehensive_history.extend(valid_output)
comprehensive_history.append(invalid_input)
comprehensive_history.append(invalid_output)
curated_history = []
curated_history.append(valid_input)
curated_history.extend(valid_output)
models_module = models.Models(mock_api_client)
chats_module = chats.Chats(modules=models_module)
chat = chats_module.create(
model='gemini-2.5-flash', history=comprehensive_history
)
assert chat.get_history() == comprehensive_history
assert chat.get_history(curated=True) == curated_history
def test_chat_with_empty_text_part(mock_generate_content_with_empty_text_part):
models_module = models.Models(mock_api_client)
chats_module = chats.Chats(modules=models_module)
chat = chats_module.create(model='gemini-2.5-flash')
chat.send_message('Hello')
expected_comprehensive_history = [
types.UserContent(parts=[types.Part.from_text(text='Hello')]),
types.Content(
parts=[types.Part(text='')],
role='model',
),
]
assert chat.get_history() == expected_comprehensive_history
assert chat.get_history(curated=True) == expected_comprehensive_history
def test_chat_with_empty_content(mock_generate_content_empty_content):
models_module = models.Models(mock_api_client)
chats_module = chats.Chats(modules=models_module)
chat = chats_module.create(model='gemini-2.5-flash')
chat.send_message('Hello')
expected_comprehensive_history = [
types.UserContent(parts=[types.Part.from_text(text='Hello')]),
types.Content(
parts=[],
role='model',
),
]
assert chat.get_history() == expected_comprehensive_history
assert not chat.get_history(curated=True)
def test_chat_stream_with_empty_text_part(
mock_generate_content_stream_with_empty_text_part,
):
models_module = models.Models(mock_api_client)
chats_module = chats.Chats(modules=models_module)
chat = chats_module.create(model='gemini-2.5-flash')
chunks = chat.send_message_stream('Hello')
for chunk in chunks:
pass
expected_comprehensive_history = [
types.UserContent(parts=[types.Part.from_text(text='Hello')]),
types.Content(
parts=[types.Part(text='')],
role='model',
),
]
assert chat.get_history() == expected_comprehensive_history
assert chat.get_history(curated=True) == expected_comprehensive_history
def test_chat_stream_with_empty_content(
mock_generate_content_stream_empty_content,
):
models_module = models.Models(mock_api_client)
chats_module = chats.Chats(modules=models_module)
chat = chats_module.create(model='gemini-2.5-flash')
chunks = chat.send_message_stream('Hello')
for chunk in chunks:
pass
expected_comprehensive_history = [
types.UserContent(parts=[types.Part.from_text(text='Hello')]),
types.Content(
parts=[],
role='model',
),
]
assert chat.get_history() == expected_comprehensive_history
assert not chat.get_history(curated=True)
def test_chat_with_afc_history(mock_generate_content_afc_history):
models_module = models.Models(mock_api_client)
chats_module = chats.Chats(modules=models_module)
chat = chats_module.create(model='gemini-2.5-flash')
chat.send_message('Hello')
expected_history = AFC_HISTORY + [
types.Content(
role='model',
parts=[types.Part.from_text(text='afc output')],
),
]
assert chat.get_history() == expected_history
assert chat.get_history(curated=True) == expected_history
def test_chat_stream_with_afc_history(mock_generate_content_stream_afc_history):
models_module = models.Models(mock_api_client)
chats_module = chats.Chats(modules=models_module)
chat = chats_module.create(model='gemini-2.5-flash')
chunks = chat.send_message_stream('Hello')
for chunk in chunks:
pass
expected_history = AFC_HISTORY + [
types.Content(
role='model',
parts=[types.Part.from_text(text='afc output')],
),
]
assert chat.get_history() == expected_history
assert chat.get_history(curated=True) == expected_history

View File

@@ -0,0 +1,941 @@
# 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 json
import os
import sys
from pydantic import BaseModel
from pydantic import ValidationError
import pytest
from .. import pytest_helper
from ... import errors
from ... import types
try:
from mcp import types as mcp_types
from mcp import ClientSession as McpClientSession
except ImportError as e:
import sys
if sys.version_info < (3, 10):
raise ImportError(
'MCP Tool requires Python 3.10 or above. Please upgrade your Python'
' version.'
) from e
else:
raise e
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
)
pytest_plugins = ('pytest_asyncio',)
MODEL_NAME = 'gemini-2.5-flash'
def divide_intergers_with_customized_math_rule(
numerator: int, denominator: int
) -> int:
"""Divides two integers with customized math rule."""
return numerator // denominator + 1
def square_integer(given_integer: int) -> int:
return given_integer*given_integer
def power_disco_ball(power: bool) -> bool:
"""Powers the spinning disco ball."""
print(f"Disco ball is {'spinning!' if power else 'stopped.'}")
return True
def start_music(energetic: bool, loud: bool, bpm: int) -> str:
"""Play some music matching the specified parameters.
Args:
energetic: Whether the music is energetic or not.
loud: Whether the music is loud or not.
bpm: The beats per minute of the music.
Returns: The name of the song being played.
"""
print(f"Starting music! {energetic=} {loud=}, {bpm=}")
return "Never gonna give you up."
def dim_lights(brightness: float) -> bool:
"""Dim the lights.
Args:
brightness: The brightness of the lights, 0.0 is off, 1.0 is full.
"""
print(f"Lights are now set to {brightness:.0%}")
return True
def test_text(client):
chat = client.chats.create(model=MODEL_NAME)
chat.send_message(
'tell me a story in 100 words',
)
def test_part(client):
chat = client.chats.create(model=MODEL_NAME)
chat.send_message(
types.Part.from_text(text='tell me a story in 100 words'),
)
def test_parts(client):
chat = client.chats.create(model=MODEL_NAME)
chat.send_message(
[
types.Part.from_text(text='tell me a US city'),
types.Part.from_text(text='the city is in west coast'),
],
)
def test_image(client, image_jpeg):
chat = client.chats.create(model=MODEL_NAME)
chat.send_message(
[
'what is the image about?',
image_jpeg,
],
)
def test_thinking_budget(client):
"""Tests that the thinking budget is respected and generates thoughts."""
chat = client.chats.create(
model=MODEL_NAME,
config={
'thinking_config': {
'include_thoughts': True,
'thinking_budget': 10000,
},
},
)
response1 = chat.send_message(
'what is the sum of natural numbers from 1 to 100?',
)
has_thought1 = False
if response1.candidates:
for candidate in response1.candidates:
for part in candidate.content.parts:
if part.thought:
has_thought1 = True
break
assert has_thought1
response2 = chat.send_message(
'can you help me to understand the logic better?'
)
has_thought2 = False
if response2.candidates:
for candidate in response2.candidates:
for part in candidate.content.parts:
if part.thought:
has_thought2 = True
break
assert has_thought2
def test_thinking_budget_stream(client):
"""Tests that the thinking budget is respected and generates thoughts."""
chat = client.chats.create(
model=MODEL_NAME,
config={
'thinking_config': {
'include_thoughts': True,
'thinking_budget': 10000,
},
},
)
has_thought1 = False
for chunk in chat.send_message_stream(
'what is the sum of natural numbers from 1 to 100?',
):
if chunk.candidates:
for candidate in chunk.candidates:
for part in candidate.content.parts:
if part.thought:
has_thought1 = True
break
assert has_thought1
has_thought2 = False
for chunk in chat.send_message_stream(
'can you help me to understand the logic better?'
):
if chunk.candidates:
for candidate in chunk.candidates:
for part in candidate.content.parts:
if part.thought:
has_thought2 = True
break
assert has_thought2
def test_google_cloud_storage_uri(client):
chat = client.chats.create(model=MODEL_NAME)
with pytest_helper.exception_if_mldev(client, errors.ClientError):
chat.send_message(
[
'what is the image about?',
types.Part.from_uri(
file_uri='gs://unified-genai-dev/imagen-inputs/google_small.png',
mime_type='image/png',
),
],
)
def test_uploaded_file_uri(client):
chat = client.chats.create(model=MODEL_NAME)
with pytest_helper.exception_if_vertex(client, errors.ClientError):
chat.send_message(
[
'what is the image about?',
types.Part.from_uri(
file_uri='https://generativelanguage.googleapis.com/v1beta/files/az606f58k7zj',
mime_type='image/png',
),
],
)
def test_config_override(client):
chat_config = {'candidate_count': 1}
chat = client.chats.create(model=MODEL_NAME, config=chat_config)
request_config = {'candidate_count': 2}
request_config_response = chat.send_message(
'tell me a story in 100 words',
config=request_config)
default_config_response = chat.send_message(
'tell me a story in 100 words')
assert len(request_config_response.candidates) == 2
assert len(default_config_response.candidates) == 1
def test_history(client):
history = [
types.Content(
role='user', parts=[types.Part.from_text(text='define a=5, b=10')]
),
types.Content(
role='model',
parts=[types.Part.from_text(text='Hello there! how can I help you?')],
),
]
chat = client.chats.create(model=MODEL_NAME, history=history)
chat.send_message('what is a + b?')
assert len(chat.get_history()) > 2
def test_send_2_messages(client):
chat = client.chats.create(model=MODEL_NAME)
chat.send_message('write a python function to check if a year is a leap year')
chat.send_message('write a unit test for the function')
def test_with_afc_history(client):
chat = client.chats.create(
model='gemini-2.0-flash-exp',
config={'tools': [divide_intergers_with_customized_math_rule]},
)
_ = chat.send_message('what is the result of 100/2?')
chat_history = chat.get_history()
assert len(chat_history) == 4
assert chat_history[0].role == 'user'
assert chat_history[0].parts[0].text == 'what is the result of 100/2?'
assert chat_history[1].role == 'model'
assert (
chat_history[1].parts[0].function_call.name
== 'divide_intergers_with_customized_math_rule'
)
assert chat_history[1].parts[0].function_call.args == {
'numerator': 100,
'denominator': 2,
}
assert chat_history[2].role == 'user'
assert (
chat_history[2].parts[0].function_response.name
== 'divide_intergers_with_customized_math_rule'
)
assert chat_history[2].parts[0].function_response.response == {'result': 51}
assert chat_history[3].role == 'model'
assert '51' in chat_history[3].parts[0].text
def test_existing_chat_history_extends_afc_history(client):
chat = client.chats.create(
model='gemini-2.0-flash-exp',
config={'tools': [divide_intergers_with_customized_math_rule]},
)
_ = chat.send_message('hello')
_ = chat.send_message('could you help me with a math problem?')
_ = chat.send_message('what is the result of 100/2?')
chat_history = chat.get_history()
content_strings = []
for content in chat_history:
content_strings.append(content.model_dump_json())
# checks that the history is not duplicated
assert len(content_strings) == len(set(content_strings))
@pytest.mark.skipif(
sys.version_info >= (3, 13),
reason=(
'object type is dumped as <Type.OBJECT: "OBJECT"> as opposed to'
' "OBJECT" in Python 3.13'
),
)
def test_with_afc_multiple_remote_calls(client):
house_fns = [power_disco_ball, start_music, dim_lights]
config = {
'tools': house_fns,
# Force the model to act (call 'any' function), instead of chatting.
'tool_config': {
'function_calling_config': {
'mode': 'ANY',
}
},
'automatic_function_calling': {
'maximum_remote_calls': 3,
}
}
chat = client.chats.create(model=MODEL_NAME, config=config)
chat.send_message('Turn this place into a party!')
curated_history = chat.get_history()
assert len(curated_history) == 8
assert curated_history[0].role == 'user'
assert curated_history[0].parts[0].text == 'Turn this place into a party!'
assert curated_history[1].role == 'model'
assert len(curated_history[1].parts) == 3
for part in curated_history[1].parts:
assert part.function_call
assert curated_history[2].role == 'user'
assert len(curated_history[2].parts) == 3
for part in curated_history[2].parts:
assert part.function_response
assert curated_history[3].role == 'model'
assert len(curated_history[3].parts) == 3
for part in curated_history[3].parts:
assert part.function_call
assert curated_history[4].role == 'user'
assert len(curated_history[4].parts) == 3
for part in curated_history[4].parts:
assert part.function_response
assert curated_history[5].role == 'model'
assert len(curated_history[5].parts) == 3
for part in curated_history[5].parts:
assert part.function_call
assert curated_history[6].role == 'user'
assert len(curated_history[6].parts) == 3
for part in curated_history[6].parts:
assert part.function_response
assert curated_history[7].role == 'model'
assert len(curated_history[7].parts) == 3
for part in curated_history[7].parts:
assert part.function_call
@pytest.mark.skipif(
sys.version_info >= (3, 13),
reason=(
'object type is dumped as <Type.OBJECT: "OBJECT"> as opposed to'
' "OBJECT" in Python 3.13'
),
)
def test_with_afc_multiple_remote_calls_async(client):
house_fns = [power_disco_ball, start_music, dim_lights]
config = {
'tools': house_fns,
# Force the model to act (call 'any' function), instead of chatting.
'tool_config': {
'function_calling_config': {
'mode': 'ANY',
}
},
'automatic_function_calling': {
'maximum_remote_calls': 3,
}
}
chat = client.chats.create(model=MODEL_NAME, config=config)
chat.send_message('Turn this place into a party!')
curated_history = chat.get_history()
assert len(curated_history) == 8
assert curated_history[0].role == 'user'
assert curated_history[0].parts[0].text == 'Turn this place into a party!'
assert curated_history[1].role == 'model'
assert len(curated_history[1].parts) == 3
for part in curated_history[1].parts:
assert part.function_call
assert curated_history[2].role == 'user'
assert len(curated_history[2].parts) == 3
for part in curated_history[2].parts:
assert part.function_response
assert curated_history[3].role == 'model'
assert len(curated_history[3].parts) == 3
for part in curated_history[3].parts:
assert part.function_call
assert curated_history[4].role == 'user'
assert len(curated_history[4].parts) == 3
for part in curated_history[4].parts:
assert part.function_response
assert curated_history[5].role == 'model'
assert len(curated_history[5].parts) == 3
for part in curated_history[5].parts:
assert part.function_call
assert curated_history[6].role == 'user'
assert len(curated_history[6].parts) == 3
for part in curated_history[6].parts:
assert part.function_response
assert curated_history[7].role == 'model'
assert len(curated_history[7].parts) == 3
for part in curated_history[7].parts:
assert part.function_call
def test_with_afc_disabled(client):
chat = client.chats.create(
model='gemini-2.0-flash-exp',
config={
'tools': [square_integer],
'automatic_function_calling': {'disable': True},
},
)
chat.send_message(
'Do the square of 3.',
)
chat_history = chat.get_history()
assert len(chat_history) == 2
assert chat_history[0].role == 'user'
assert chat_history[0].parts[0].text == 'Do the square of 3.'
assert chat_history[1].role == 'model'
assert chat_history[1].parts[0].function_call.name == 'square_integer'
assert chat_history[1].parts[0].function_call.args == {
'given_integer': 3,
}
@pytest.mark.asyncio
async def test_with_afc_history_async(client):
chat = client.aio.chats.create(
model='gemini-2.0-flash-exp',
config={'tools': [divide_intergers_with_customized_math_rule]},
)
_ = await chat.send_message('what is the result of 100/2?')
chat_history = chat.get_history()
assert len(chat_history) == 4
assert chat_history[0].role == 'user'
assert chat_history[0].parts[0].text == 'what is the result of 100/2?'
assert chat_history[1].role == 'model'
assert (
chat_history[1].parts[0].function_call.name
== 'divide_intergers_with_customized_math_rule'
)
assert chat_history[1].parts[0].function_call.args == {
'numerator': 100,
'denominator': 2,
}
assert chat_history[2].role == 'user'
assert (
chat_history[2].parts[0].function_response.name
== 'divide_intergers_with_customized_math_rule'
)
assert chat_history[2].parts[0].function_response.response == {'result': 51}
assert chat_history[3].role == 'model'
assert '51' in chat_history[3].parts[0].text
@pytest.mark.asyncio
async def test_with_afc_disabled_async(client):
chat = client.aio.chats.create(
model='gemini-2.0-flash-exp',
config={
'tools': [square_integer],
'automatic_function_calling': {'disable': True},
},
)
await chat.send_message(
'Do the square of 3.',
)
chat_history = chat.get_history()
assert len(chat_history) == 2
assert chat_history[0].role == 'user'
assert chat_history[0].parts[0].text == 'Do the square of 3.'
assert chat_history[1].role == 'model'
assert chat_history[1].parts[0].function_call.name == 'square_integer'
assert chat_history[1].parts[0].function_call.args == {
'given_integer': 3,
}
def test_stream_text(client):
chat = client.chats.create(model=MODEL_NAME)
chunks = 0
for chunk in chat.send_message_stream(
'tell me a story in 100 words',
):
chunks += 1
assert chunks > 1
def test_stream_part(client):
chat = client.chats.create(model=MODEL_NAME)
chunks = 0
for chunk in chat.send_message_stream(
types.Part.from_text(text='tell me a story in 100 words'),
):
chunks += 1
assert chunks > 1
def test_stream_parts(client):
chat = client.chats.create(model=MODEL_NAME)
chunks = 0
for chunk in chat.send_message_stream(
[
types.Part.from_text(text='tell me a story in 100 words'),
types.Part.from_text(text='the story is about a car'),
],
):
chunks += 1
assert chunks > 2
def test_stream_config_override(client):
chat_config = {'response_mime_type': 'text/plain'}
chat = client.chats.create(model=MODEL_NAME, config=chat_config)
request_config = {'response_mime_type': 'application/json'}
request_config_text = ''
for chunk in chat.send_message_stream(
'tell me a story in 100 words', config=request_config
):
request_config_text += chunk.text
default_config_text = ''
for chunk in chat.send_message_stream('tell me a story in 100 words'):
default_config_text += chunk.text
assert json.loads(request_config_text)
with pytest.raises(json.JSONDecodeError):
json.loads(default_config_text)
def test_stream_function_calling(client):
chat = client.chats.create(
model='gemini-2.0-flash-exp',
config={'tools': [divide_intergers_with_customized_math_rule]},
)
# Now we support AFC.
for chunk in chat.send_message_stream(
'what is the result of 100/2?',
):
pass
for chunk in chat.send_message_stream(
'what is the result of 50/2?',
):
pass
chat_history = chat.get_history()
assert chat_history[0].role == 'user'
assert chat_history[0].parts[0].text == 'what is the result of 100/2?'
assert chat_history[1].role == 'model'
assert (
chat_history[1].parts[0].function_call.name
== 'divide_intergers_with_customized_math_rule'
)
assert chat_history[1].parts[0].function_call.args == {
'numerator': 100,
'denominator': 2,
}
def test_stream_send_2_messages(client):
chat = client.chats.create(model=MODEL_NAME)
for chunk in chat.send_message_stream(
'write a python function to check if a year is a leap year'
):
pass
for chunk in chat.send_message_stream('write a unit test for the function'):
pass
@pytest.mark.asyncio
async def test_async_text(client):
chat = client.aio.chats.create(model=MODEL_NAME)
await chat.send_message('tell me a story in 100 words')
@pytest.mark.asyncio
async def test_async_part(client):
chat = client.aio.chats.create(model=MODEL_NAME)
await chat.send_message(types.Part.from_text(text='tell me a story in 100 words'))
@pytest.mark.asyncio
async def test_async_parts(client):
chat = client.aio.chats.create(model=MODEL_NAME)
await chat.send_message(
[
types.Part.from_text(text='tell me a US city'),
types.Part.from_text(text='the city is in west coast'),
],
)
@pytest.mark.asyncio
async def test_async_config_override(client):
chat_config = {'candidate_count': 1}
chat = client.aio.chats.create(model=MODEL_NAME, config=chat_config)
request_config = {'candidate_count': 2}
request_config_response = await chat.send_message(
'tell me a story in 100 words',
config=request_config)
default_config_response = await chat.send_message(
'tell me a story in 100 words')
assert len(request_config_response.candidates) == 2
assert len(default_config_response.candidates) == 1
@pytest.mark.asyncio
async def test_async_history(client):
history = [
types.Content(
role='user', parts=[types.Part.from_text(text='define a=5, b=10')]
),
types.Content(
role='model',
parts=[types.Part.from_text(text='Hello there! how can I help you?')],
),
]
chat = client.aio.chats.create(model=MODEL_NAME, history=history)
await chat.send_message('what is a + b?')
assert len(chat.get_history()) > 2
@pytest.mark.asyncio
async def test_async_stream_text(client):
chat = client.aio.chats.create(model=MODEL_NAME)
chunks = 0
async for chunk in await chat.send_message_stream('tell me a story in 100 words'):
chunks += 1
assert chunks > 1
@pytest.mark.asyncio
async def test_async_stream_part(client):
chat = client.aio.chats.create(model=MODEL_NAME)
chunks = 0
async for chunk in await chat.send_message_stream(
types.Part.from_text(text='tell me a story in 100 words')
):
chunks += 1
assert chunks > 1
@pytest.mark.asyncio
async def test_async_stream_parts(client):
chat = client.aio.chats.create(model=MODEL_NAME)
chunks = 0
async for chunk in await chat.send_message_stream(
[
types.Part.from_text(text='tell me a story in 100 words'),
types.Part.from_text(text='the story is about a car'),
],
):
chunks += 1
assert chunks > 1
@pytest.mark.asyncio
async def test_async_stream_config_override(client):
chat_config = {'response_mime_type': 'text/plain'}
chat = client.aio.chats.create(model=MODEL_NAME, config=chat_config)
request_config = {'response_mime_type': 'application/json'}
request_config_text = ''
async for chunk in await chat.send_message_stream(
'tell me a story in 100 words', config=request_config
):
request_config_text += chunk.text
default_config_text = ''
async for chunk in await chat.send_message_stream('tell me family friendly story in 100 words'):
default_config_text += chunk.text
assert json.loads(request_config_text)
with pytest_helper.exception_if_mldev(client, json.JSONDecodeError):
json.loads(default_config_text)
@pytest.mark.asyncio
async def test_async_stream_function_calling(client):
chat = client.aio.chats.create(
model='gemini-2.0-flash-exp',
config={'tools': [divide_intergers_with_customized_math_rule]},
)
# Now we support AFC.
async for chunk in await chat.send_message_stream('what is the result of 100/2?'):
pass
async for chunk in await chat.send_message_stream('what is the result of 50/2?'):
pass
chat_history = chat.get_history()
assert chat_history[0].role == 'user'
assert chat_history[0].parts[0].text == 'what is the result of 100/2?'
assert chat_history[1].role == 'model'
assert (
chat_history[1].parts[0].function_call.name
== 'divide_intergers_with_customized_math_rule'
)
assert chat_history[1].parts[0].function_call.args == {
'numerator': 100,
'denominator': 2,
}
@pytest.mark.asyncio
async def test_async_stream_send_2_messages(client):
chat = client.aio.chats.create(model=MODEL_NAME)
async for chunk in await chat.send_message_stream(
'write a python function to check if a year is a leap year'
):
pass
async for chunk in await chat.send_message_stream(
'write a unit test for the function'
):
pass
def test_mcp_tools(client):
chat = client.chats.create(
model='gemini-2.0-flash-exp',
config={'tools': [
mcp_types.Tool(
name='get_weather',
description='Get the weather in a city.',
inputSchema={
'type': 'object',
'properties': {'location': {'type': 'string'}},
},
)
],},
)
response = chat.send_message('What is the weather in Boston?')
response = chat.send_message('What is the weather in San Francisco?')
def test_mcp_tools_stream(client):
chat = client.chats.create(
model='gemini-2.0-flash-exp',
config={'tools': [
mcp_types.Tool(
name='get_weather',
description='Get the weather in a city.',
inputSchema={
'type': 'object',
'properties': {'location': {'type': 'string'}},
},
)
],
},
)
for chunk in chat.send_message_stream(
'What is the weather in Boston?'
):
pass
for chunk in chat.send_message_stream(
'What is the weather in San Francisco?'
):
pass
@pytest.mark.asyncio
async def test_async_mcp_tools(client):
chat = client.aio.chats.create(
model='gemini-2.0-flash-exp',
config={'tools': [
mcp_types.Tool(
name='get_weather',
description='Get the weather in a city.',
inputSchema={
'type': 'object',
'properties': {'location': {'type': 'string'}},
},
)
],},
)
await chat.send_message('What is the weather in Boston?');
await chat.send_message('What is the weather in San Francisco?');
@pytest.mark.asyncio
async def test_async_mcp_tools_stream(client):
chat = client.aio.chats.create(
model='gemini-2.0-flash-exp',
config={'tools': [
mcp_types.Tool(
name='get_weather',
description='Get the weather in a city.',
inputSchema={
'type': 'object',
'properties': {'location': {'type': 'string'}},
},
)
],
},
)
async for chunk in await chat.send_message_stream(
'What is the weather in Boston?'
):
pass
async for chunk in await chat.send_message_stream(
'What is the weather in San Francisco?'
):
pass
def test_server_side_mcp_tools(client):
with pytest_helper.exception_if_vertex(client, ValueError):
chat = client.chats.create(
model='gemini-2.5-flash',
config={
'tools': [
{
'mcp_servers': [
{
'name': 'weather_server',
'streamable_http_transport': {
'url': (
'https://gemini-api-demos.uc.r.appspot.com/mcp'
),
'headers': {
'AUTHORIZATION': 'Bearer github_pat_XXXX',
},
'timeout': '10s',
},
},
],
},
],
},
)
response = chat.send_message('What is the weather in Boston on 02/02/2026?')
response = chat.send_message(
'What is the weather in San Francisco on 02/02/2026?'
)
def test_server_side_mcp_tools_stream(client):
with pytest_helper.exception_if_vertex(client, ValueError):
chat = client.chats.create(
model='gemini-2.5-flash',
config={
'tools': [
{
'mcp_servers': [
{
'name': 'weather_server',
'streamable_http_transport': {
'url': (
'https://gemini-api-demos.uc.r.appspot.com/mcp'
),
'headers': {
'AUTHORIZATION': 'Bearer github_pat_XXXX',
},
'timeout': '10s',
},
},
],
},
],
},
)
for chunk in chat.send_message_stream(
'What is the weather in Boston on 02/02/2026?'
):
pass
for chunk in chat.send_message_stream(
'What is the weather in San Francisco on 02/02/2026?'
):
pass
@pytest.mark.asyncio
async def test_async_server_side_mcp_tools(client):
with pytest_helper.exception_if_vertex(client, ValueError):
chat = client.aio.chats.create(
model='gemini-2.5-flash',
config={
'tools': [
{
'mcp_servers': [
{
'name': 'weather_server',
'streamable_http_transport': {
'url': (
'https://gemini-api-demos.uc.r.appspot.com/mcp'
),
'headers': {
'AUTHORIZATION': 'Bearer github_pat_XXXX',
},
'timeout': '10s',
},
},
],
},
],
},
)
await chat.send_message('What is the weather in Boston on 02/02/2026?')
await chat.send_message(
'What is the weather in San Francisco on 02/02/2026?'
)

View File

@@ -0,0 +1,90 @@
# 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.
#
from ... import types
from ...chats import _validate_response
def test_validate_response_default_response():
response = types.GenerateContentResponse()
assert not _validate_response(response)
def test_validate_response_empty_content():
response = types.GenerateContentResponse(candidates=[])
assert not _validate_response(response)
def test_validate_response_empty_parts():
response = types.GenerateContentResponse(
candidates=[types.Candidate(content=types.Content(parts=[]))]
)
assert not _validate_response(response)
def test_validate_response_empty_part():
response = types.GenerateContentResponse(
candidates=[types.Candidate(content=types.Content(parts=[types.Part()]))]
)
assert not _validate_response(response)
def test_validate_response_part_with_empty_text():
response = types.GenerateContentResponse(
candidates=[
types.Candidate(content=types.Content(parts=[types.Part(text='')]))
]
)
assert _validate_response(response)
def test_validate_response_part_with_text():
response = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
parts=[types.Part(text='response from model')]
)
)
]
)
assert _validate_response(response)
def test_validate_response_part_with_function_call():
response = types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
parts=[
types.Part(
function_call=types.FunctionCall(
name='foo', args={'bar': 'baz'}
)
)
]
)
)
]
)
assert _validate_response(response)

View File

@@ -0,0 +1,17 @@
# 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.
#
"""Tests for the Google GenAI SDK."""

View File

@@ -0,0 +1,427 @@
# 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.
#
"""Tests for async stream."""
import asyncio
from typing import List
from unittest import mock
from unittest.mock import AsyncMock
from unittest.mock import MagicMock
from unittest.mock import patch
import pytest
try:
import aiohttp
AIOHTTP_NOT_INSTALLED = False
except ImportError:
AIOHTTP_NOT_INSTALLED = True
aiohttp = mock.MagicMock()
import httpx
from ... import _api_client as api_client
from ... import Client
from ... import errors
from ... import types
EVENT_STREAM_DATA_WITH_ERROR = [
b'{"candidates":[{"content":{"parts":[{"text":"test"}],"role":"model"}}]}',
b'\n',
b'{"error":{"code":500,"message":"Error","status":"INTERNAL"}}',
]
class MockHTTPXResponse(httpx.Response):
"""Mock httpx.Response class for testing."""
def __init__(self, lines: List[str]):
self.aiter_lines = MagicMock()
self.aiter_lines.return_value.__aiter__ = MagicMock(
return_value=self._async_line_iterator(lines)
)
self.aclose = AsyncMock()
async def _async_line_iterator(self, lines: List[str]):
for line in lines:
yield line
class MockAIOHTTPResponse(aiohttp.ClientResponse):
def __init__(self, lines: List[str]):
self.content = MagicMock()
self.content.readline = AsyncMock()
# Simulate reading lines, each ending with newline bytes for readline behavior
self._read_data = b"\n".join(line.encode("utf-8") for line in lines) + b"\n"
self._read_pos = 0
self.content.readline.side_effect = self._async_read_line
self.release = MagicMock()
async def _async_read_line(self) -> bytes:
if self._read_pos >= len(self._read_data):
return b"" # End of stream
newline_pos = self._read_data.find(b"\n", self._read_pos)
if newline_pos == -1: # Should not happen with the appended '\n'
line = self._read_data[self._read_pos :]
self._read_pos = len(self._read_data)
return line
else:
line = self._read_data[self._read_pos : newline_pos + 1]
self._read_pos = newline_pos + 1
return line
@pytest.fixture
def responses() -> api_client.HttpResponse:
return api_client.HttpResponse(headers={})
requires_aiohttp = pytest.mark.skipif(
AIOHTTP_NOT_INSTALLED, reason="aiohttp is not installed, skipping test."
)
@pytest.fixture(autouse=True)
def reset_has_aiohttp():
yield
api_client.has_aiohttp = False
def test_invalid_response_stream_type(responses: api_client.HttpResponse):
"""Tests that an invalid response stream type raises an error."""
api_client.has_aiohttp = False
with pytest.raises(
TypeError,
match=(
"Expected self.response_stream to be an httpx.Response or"
" aiohttp.ClientResponse object"
),
):
async def run():
async for _ in responses._aiter_response_stream():
pass
asyncio.run(run())
@pytest.mark.asyncio
async def test_httpx_simple_lines(responses: api_client.HttpResponse):
lines = ["hello", "world", "testing"]
mock_response = MockHTTPXResponse(lines)
responses.response_stream = mock_response
results = [line async for line in responses._aiter_response_stream()]
assert results == lines
mock_response.aiter_lines.assert_called_once()
mock_response.aclose.assert_called_once()
@pytest.mark.asyncio
async def test_httpx_data_prefix(responses: api_client.HttpResponse):
lines = ["data: { 'message': 'hello' }", "data: { 'status': 'ok' }"]
mock_response = MockHTTPXResponse(lines)
responses.response_stream = mock_response
results = [line async for line in responses._aiter_response_stream()]
assert results == ["{ 'message': 'hello' }", "{ 'status': 'ok' }"]
mock_response.aiter_lines.assert_called_once()
mock_response.aclose.assert_called_once()
@pytest.mark.asyncio
async def test_httpx_multiple_json_chunk(responses: api_client.HttpResponse):
lines = [
'{ "id": 1 }',
"",
'data: { "id": 2 }',
'data: { "id": 3 }',
]
mock_response = MockHTTPXResponse(lines)
responses.response_stream = mock_response
results = [line async for line in responses._aiter_response_stream()]
assert results == ['{ "id": 1 }', '{ "id": 2 }', '{ "id": 3 }']
mock_response.aiter_lines.assert_called_once()
mock_response.aclose.assert_called_once()
@pytest.mark.asyncio
async def test_httpx_incomplete_json_at_end(responses: api_client.HttpResponse):
lines = ['{ "partial": "data"'] # Missing closing brace
mock_response = MockHTTPXResponse(lines)
responses.response_stream = mock_response
results = [line async for line in responses._aiter_response_stream()]
# The remaining chunk is yielded
assert results == ['{ "partial": "data"']
mock_response.aiter_lines.assert_called_once()
mock_response.aclose.assert_called_once()
@pytest.mark.asyncio
async def test_httpx_empty_stream(responses: api_client.HttpResponse):
lines: List[str] = []
mock_response = MockHTTPXResponse(lines)
responses.response_stream = mock_response
results = [line async for line in responses._aiter_response_stream()]
assert results == []
mock_response.aiter_lines.assert_called_once()
mock_response.aclose.assert_called_once()
# Async aiohttp
@requires_aiohttp
@pytest.mark.asyncio
async def test_aiohttp_simple_lines(responses: api_client.HttpResponse):
api_client.has_aiohttp = True # Force aiohttp
lines = ["hello", "world", "testing"]
# Use the mock class that pretends to be aiohttp.ClientResponse
mock_response = MockAIOHTTPResponse(lines)
responses.response_stream = mock_response
results = [line async for line in responses._aiter_response_stream()]
assert results == lines
mock_response.content.readline.assert_any_call()
mock_response.release.assert_called_once()
@requires_aiohttp
@pytest.mark.asyncio
async def test_aiohttp_data_prefix(responses: api_client.HttpResponse):
api_client.has_aiohttp = True # Force aiohttp
lines = ["data: { 'message': 'hello' }", "data: { 'status': 'ok' }"]
# Use the mock class that pretends to be aiohttp.ClientResponse
mock_response = MockAIOHTTPResponse(lines)
responses.response_stream = mock_response
results = [line async for line in responses._aiter_response_stream()]
assert results == ["{ 'message': 'hello' }", "{ 'status': 'ok' }"]
mock_response.content.readline.assert_any_call()
mock_response.release.assert_called_once()
@requires_aiohttp
@pytest.mark.asyncio
async def test_aiohttp_multiple_json_chunks(responses: api_client.HttpResponse):
api_client.has_aiohttp = True # Force aiohttp
lines = [
'{ "id": 1 }',
"", # empty line to check robustness
'data: { "id": 2 }',
'data: { "id": 3 }',
]
# Use the mock class that pretends to be aiohttp.ClientResponse
mock_response = MockAIOHTTPResponse(lines)
responses.response_stream = mock_response
results = [line async for line in responses._aiter_response_stream()]
assert results == ['{ "id": 1 }', '{ "id": 2 }', '{ "id": 3 }']
mock_response.content.readline.assert_any_call()
mock_response.release.assert_called_once()
@requires_aiohttp
@pytest.mark.asyncio
async def test_aiohttp_incomplete_json_at_end(
responses: api_client.HttpResponse,
):
api_client.has_aiohttp = True # Force aiohttp
lines = ['{ "partial": "data"'] # Missing closing brace
# Use the mock class that pretends to be aiohttp.ClientResponse
mock_response = MockAIOHTTPResponse(lines)
responses.response_stream = mock_response
results = [line async for line in responses._aiter_response_stream()]
assert results == ['{ "partial": "data"']
mock_response.content.readline.assert_any_call()
mock_response.release.assert_called_once()
def mock_response(chunks):
mock_stream = MagicMock(spec=httpx.SyncByteStream)
mock_stream.__iter__.return_value = chunks
return httpx.Response(
status_code=200,
stream=mock_stream,
)
@patch('httpx.Client.send')
def test_error_event_in_streamed_responses_bad_json(mock_send_method):
with_bad_json = [
b'{"candidates":[{"content":{"parts":[{"text":"test"}],"role":"model"}}]}',
b'\n',
b'{"error": bad_json}',
]
mock_send_method.return_value = mock_response(with_bad_json)
client = api_client.BaseApiClient(api_key='test_api_key')
stream = client.request_streamed('POST', 'models/gemini-2.5-flash', {})
chunk = next(stream)
assert chunk == types.HttpResponse(
headers={},
body=(
'{"candidates": [{"content": {"parts": [{"text": "test"}], "role":'
' "model"}}]}'
),
)
with pytest.raises(errors.UnknownApiResponseError):
next(stream)
@patch('httpx.Client.send')
def test_error_event_in_streamed_responses(mock_send_method):
mock_send_method.return_value = mock_response(EVENT_STREAM_DATA_WITH_ERROR)
client = api_client.BaseApiClient(api_key='test_api_key')
stream = client.request_streamed('POST', 'models/gemini-2.5-flash', {})
chunk = next(stream)
assert chunk == types.HttpResponse(
body=(
'{"candidates": [{"content": {"parts": [{"text": "test"}], "role":'
' "model"}}]}'
),
headers={},
)
with pytest.raises(errors.ServerError):
next(stream)
@patch('httpx.Client.send')
def test_error_event_in_generate_content_stream(mock_send_method):
mock_send_method.return_value = mock_response(EVENT_STREAM_DATA_WITH_ERROR)
client = Client(api_key='test_api_key')
generated_response = client.models.generate_content_stream(
model='gemini-2.5-flash',
contents='Tell me a story in 300 words.',
)
chunk = next(generated_response)
assert chunk == types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
parts=[
types.Part(
text='test'
),
],
role='model'
)
),
],
sdk_http_response=types.HttpResponse(headers={})
)
with pytest.raises(errors.ServerError):
next(generated_response)
async def _async_httpx_response(_):
mock_stream = MagicMock(spec=httpx.AsyncByteStream)
mock_stream.__aiter__.return_value = EVENT_STREAM_DATA_WITH_ERROR
mock_stream.aclose = AsyncMock()
return httpx.Response(
status_code=200,
stream=mock_stream,
)
@patch('httpx.AsyncBaseTransport')
@pytest.mark.asyncio
async def test_error_event_in_streamed_responses_async(mock_transport):
client = api_client.BaseApiClient(
api_key='test_api_key',
http_options=types.HttpOptions(
async_client_args={'transport': mock_transport}
),
)
mock_transport.handle_async_request = _async_httpx_response
mock_transport.aclose = AsyncMock()
resp = await client.async_request_streamed(
'POST', 'models/gemini-2.5-flash', {'key': 'value'}
)
chunk = await anext(resp)
assert chunk == types.HttpResponse(
headers={},
body=(
'{"candidates": [{"content": {"parts": [{"text": "test"}], "role":'
' "model"}}]}'
),
)
with pytest.raises(errors.ServerError):
await anext(resp)
@patch('httpx.AsyncBaseTransport')
@pytest.mark.asyncio
async def test_error_event_in_generate_content_stream_async(mock_transport):
client = Client(
api_key='test_api_key',
http_options=types.HttpOptions(
async_client_args={'transport': mock_transport}
),
)
mock_transport.handle_async_request = _async_httpx_response
mock_transport.aclose = AsyncMock()
generated_response = await client.aio.models.generate_content_stream(
model='gemini-2.5-flash',
contents='Tell me a story in 300 words.',
)
chunk = await anext(generated_response)
assert chunk == types.GenerateContentResponse(
candidates=[
types.Candidate(
content=types.Content(
parts=[
types.Part(
text='test'
),
],
role='model'
)
),
],
sdk_http_response=types.HttpResponse(headers={})
)
with pytest.raises(errors.ServerError):
await anext(generated_response)

View File

@@ -0,0 +1,204 @@
# 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.
#
"""Tests for closing the clients and context managers."""
import asyncio
from unittest import mock
from google.oauth2 import credentials
import pytest
try:
import aiohttp
AIOHTTP_NOT_INSTALLED = False
except ImportError:
AIOHTTP_NOT_INSTALLED = True
aiohttp = mock.MagicMock()
from ... import _api_client as api_client
from ... import Client
requires_aiohttp = pytest.mark.skipif(
AIOHTTP_NOT_INSTALLED, reason='aiohttp is not installed, skipping test.'
)
def test_close_httpx_client():
"""Tests that the httpx client is closed when the client is closed."""
api_client.has_aiohttp = False
client = Client(
vertexai=True,
project='test_project',
location='global',
http_options=api_client.HttpOptions(client_args={'max_redirects': 10}),
)
client.close()
assert client._api_client._httpx_client.is_closed
def test_httpx_client_context_manager():
"""Tests that the httpx client is closed when the client is closed."""
api_client.has_aiohttp = False
with Client(
vertexai=True,
project='test_project',
location='global',
http_options=api_client.HttpOptions(client_args={'max_redirects': 10}),
) as client:
pass
assert not client._api_client._httpx_client.is_closed
assert client._api_client._httpx_client.is_closed
@pytest.mark.asyncio
async def test_aclose_httpx_client():
"""Tests that the httpx async client is closed when the client is closed."""
api_client.has_aiohttp = False
async_client = Client(
vertexai=True,
project='test_project',
location='global',
).aio
await async_client.aclose()
assert async_client._api_client._async_httpx_client.is_closed
@pytest.mark.asyncio
async def test_async_httpx_client_context_manager():
"""Tests that the httpx async client is closed when the client is closed."""
api_client.has_aiohttp = False
async with Client(
vertexai=True,
project='test_project',
location='global',
).aio as async_client:
pass
assert not async_client._api_client._async_httpx_client.is_closed
assert async_client._api_client._async_httpx_client.is_closed
@pytest.fixture
def mock_request():
mock_aiohttp_response = mock.Mock(spec=aiohttp.ClientSession.request)
mock_aiohttp_response.return_value = mock_aiohttp_response
yield mock_aiohttp_response
def _patch_auth_default():
return mock.patch(
'google.auth.default',
return_value=(credentials.Credentials('magic_token'), 'test_project'),
autospec=True,
)
async def _aiohttp_async_response(status: int):
"""Has to return a coroutine hence async."""
response = mock.Mock(spec=aiohttp.ClientResponse)
response.status = status
response.headers = {'status-code': str(status)}
response.json.return_value = {}
response.text.return_value = 'test'
return response
@requires_aiohttp
@mock.patch.object(aiohttp.ClientSession, 'request', autospec=True)
def test_aclose_aiohttp_session(mock_request):
"""Tests that the aiohttp session is closed when the client is closed."""
api_client.has_aiohttp = True
async def run():
mock_request.side_effect = (
aiohttp.ClientConnectorError(
connection_key=aiohttp.client_reqrep.ConnectionKey(
'localhost', 80, False, True, None, None, None
),
os_error=OSError,
),
_aiohttp_async_response(200),
)
with _patch_auth_default():
async_client = Client(
vertexai=True,
project='test_project',
location='global',
http_options=api_client.HttpOptions(
async_client_args={'trust_env': False}
),
).aio
# aiohttp session is created in the first request instead of client
# initialization.
_ = await async_client._api_client._async_request_once(
api_client.HttpRequest(
method='GET',
url='https://example.com',
headers={},
data=None,
timeout=None,
)
)
assert async_client._api_client._aiohttp_session is not None
assert not async_client._api_client._aiohttp_session.closed
# Close the client and check that the session is closed.
await async_client.aclose()
assert async_client._api_client._aiohttp_session.closed
asyncio.run(run())
@requires_aiohttp
@mock.patch.object(aiohttp.ClientSession, 'request', autospec=True)
def test_aiohttp_session_context_manager(mock_request):
"""Tests that the aiohttp session is closed when the client is closed."""
api_client.has_aiohttp = True
async def run():
mock_request.side_effect = (
aiohttp.ClientConnectorError(
connection_key=aiohttp.client_reqrep.ConnectionKey(
'localhost', 80, False, True, None, None, None
),
os_error=OSError,
),
_aiohttp_async_response(200),
)
with _patch_auth_default():
async with Client(
vertexai=True,
project='test_project',
location='global',
http_options=api_client.HttpOptions(
async_client_args={'trust_env': False}
),
).aio as async_client:
# aiohttp session is created in the first request instead of client
# initialization.
_ = await async_client._api_client._async_request_once(
api_client.HttpRequest(
method='GET',
url='https://example.com',
headers={},
data=None,
timeout=None,
)
)
assert async_client._api_client._aiohttp_session is not None
assert not async_client._api_client._aiohttp_session.closed
assert async_client._api_client._aiohttp_session.closed
asyncio.run(run())

View File

@@ -0,0 +1,221 @@
# 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.
#
"""Tests for client behavior when issuing requests."""
from ... import _api_client as api_client
from ... import Client
from ... import types
def build_test_client(monkeypatch):
monkeypatch.setenv('GOOGLE_API_KEY', 'google_api_key')
return Client()
def test_join_url_path_base_url_with_trailing_slash_and_path_with_leading_slash():
base_url = 'https://fake-url.com/some_path/'
path = '/v1beta/models'
assert (
api_client.join_url_path(base_url, path)
== 'https://fake-url.com/some_path/v1beta/models'
)
def test_join_url_path_with_base_url_with_trailing_slash_and_path_without_leading_slash():
base_url = 'https://fake-url.com/some_path/'
path = 'v1beta/models'
assert (
api_client.join_url_path(base_url, path)
== 'https://fake-url.com/some_path/v1beta/models'
)
def test_join_url_path_with_base_url_without_trailing_slash_and_path_with_leading_slash():
base_url = 'https://fake-url.com/some_path'
path = '/v1beta/models'
assert (
api_client.join_url_path(base_url, path)
== 'https://fake-url.com/some_path/v1beta/models'
)
def test_join_url_path_with_base_url_without_trailing_slash_and_path_without_leading_slash():
base_url = 'https://fake-url.com/some_path'
path = 'v1beta/models'
assert (
api_client.join_url_path(base_url, path)
== 'https://fake-url.com/some_path/v1beta/models'
)
def test_join_url_path_base_url_without_path_with_trailing_slash():
base_url = 'https://fake-url.com/'
path = 'v1beta/models'
assert (
api_client.join_url_path(base_url, path)
== 'https://fake-url.com/v1beta/models'
)
def test_join_url_path_base_url_without_path_without_trailing_slash():
base_url = 'https://fake-url.com'
path = 'v1beta/models'
assert (
api_client.join_url_path(base_url, path)
== 'https://fake-url.com/v1beta/models'
)
def test_build_request_sets_library_version_headers(monkeypatch):
request_client = build_test_client(monkeypatch).models._api_client
request = request_client._build_request('GET', 'test/path', {'key': 'value'})
assert 'google-genai-sdk/' in request.headers['user-agent']
assert 'gl-python/' in request.headers['user-agent']
assert 'google-genai-sdk/' in request.headers['x-goog-api-client']
assert 'gl-python/' in request.headers['x-goog-api-client']
def test_build_request_appends_to_user_agent_headers(monkeypatch):
request_client = build_test_client(monkeypatch).models._api_client
request = request_client._build_request(
'GET',
'test/path',
{'key': 'value'},
types.HttpOptionsDict(
base_url='test/url',
api_version='1',
headers={'user-agent': 'test-user-agent'},
),
)
assert 'test-user-agent' in request.headers['user-agent']
assert 'google-genai-sdk/' in request.headers['user-agent']
assert 'gl-python/' in request.headers['user-agent']
assert 'google-genai-sdk/' in request.headers['x-goog-api-client']
def test_build_request_appends_to_goog_api_client_headers(monkeypatch):
request_client = build_test_client(monkeypatch).models._api_client
request = request_client._build_request(
'GET',
'test/path',
{'key': 'value'},
types.HttpOptionsDict(
base_url='test/url',
api_version='1',
headers={'x-goog-api-client': 'test-goog-api-client'},
),
)
assert 'google-genai-sdk/' in request.headers['user-agent']
assert 'test-goog-api-client' in request.headers['x-goog-api-client']
assert 'google-genai-sdk/' in request.headers['x-goog-api-client']
assert 'gl-python/' in request.headers['x-goog-api-client']
def test_build_request_keeps_sdk_version_headers(monkeypatch):
headers_to_inject = {}
api_client.append_library_version_headers(headers_to_inject)
assert 'google-genai-sdk/' in headers_to_inject['user-agent']
request_client = build_test_client(monkeypatch).models._api_client
request = request_client._build_request(
'GET',
'test/path',
{'key': 'value'},
types.HttpOptionsDict(
base_url='test/url',
api_version='1',
headers=headers_to_inject,
),
)
assert 'google-genai-sdk/' in request.headers['user-agent']
assert 'gl-python/' in request.headers['x-goog-api-client']
assert 'google-genai-sdk/' in request.headers['x-goog-api-client']
assert 'gl-python/' in request.headers['x-goog-api-client']
def test_build_request_with_resource_scope(monkeypatch):
monkeypatch.delenv('GOOGLE_API_KEY', raising=False)
monkeypatch.delenv('GEMINI_API_KEY', raising=False)
monkeypatch.delenv('GOOGLE_CLOUD_PROJECT', raising=False)
monkeypatch.delenv('GOOGLE_CLOUD_LOCATION', raising=False)
client = Client(
vertexai=True,
http_options=types.HttpOptionsDict(
base_url='https://custom-base-url.com',
base_url_resource_scope=types.ResourceScope.COLLECTION,
),
)
request = client.models._api_client._build_request(
'post',
'publishers/google/models/gemini-3-pro-preview',
{'key': 'value'},
)
assert request.url == 'https://custom-base-url.com/publishers/google/models/gemini-3-pro-preview'
def test_build_request_with_resource_scope_with_project_and_location(
monkeypatch,
):
monkeypatch.delenv('GOOGLE_API_KEY', raising=False)
monkeypatch.delenv('GEMINI_API_KEY', raising=False)
monkeypatch.delenv('GOOGLE_CLOUD_PROJECT', raising=False)
monkeypatch.delenv('GOOGLE_CLOUD_LOCATION', raising=False)
client = Client(
vertexai=True,
project='test-project',
location='test-location',
http_options=types.HttpOptionsDict(
base_url='https://custom-base-url.com',
base_url_resource_scope=types.ResourceScope.COLLECTION,
),
)
request = client.models._api_client._build_request(
'post',
'publishers/google/models/gemini-3-pro-preview',
{'key': 'value'},
)
assert request.url == 'https://custom-base-url.com/publishers/google/models/gemini-3-pro-preview'
def build_test_client_no_env_vars(monkeypatch):
monkeypatch.delenv('GOOGLE_API_KEY', raising=False)
monkeypatch.delenv('GEMINI_API_KEY', raising=False)
monkeypatch.delenv('GOOGLE_CLOUD_PROJECT', raising=False)
monkeypatch.delenv('GOOGLE_CLOUD_LOCATION', raising=False)
return Client(
vertexai=True,
http_options=types.HttpOptionsDict(
base_url='https://custom-base-url.com',
headers={'Authorization': 'Bearer fake_access_token'},
),
)
def test_build_request_with_custom_base_url_no_env_vars(monkeypatch):
request_client = (
build_test_client_no_env_vars(monkeypatch).models._api_client
)
request = request_client._build_request(
'GET',
'test/path',
{'key': 'value'},
)
assert request.url == 'https://custom-base-url.com'

View File

@@ -0,0 +1,104 @@
# 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.
#
"""Tests for custom clients."""
import asyncio
from unittest import mock
from google.oauth2 import credentials
import httpx
import pytest
from ... import _api_client as api_client
from ... import Client
try:
import aiohttp
AIOHTTP_NOT_INSTALLED = False
except ImportError:
AIOHTTP_NOT_INSTALLED = True
aiohttp = mock.MagicMock()
requires_aiohttp = pytest.mark.skipif(
AIOHTTP_NOT_INSTALLED, reason='aiohttp is not installed, skipping test.'
)
# Httpx
def test_constructor_with_httpx_clients():
mldev_http_options = {
'httpx_client': httpx.Client(trust_env=False),
'httpx_async_client': httpx.AsyncClient(trust_env=False),
}
vertexai_http_options = {
'httpx_client': httpx.Client(trust_env=False),
'httpx_async_client': httpx.AsyncClient(trust_env=False),
}
# Even if aiohttp is installed, expect it to be disabled when httpx clients
# are provided.
api_client.has_aiohttp = True
mldev_client = Client(
api_key='google_api_key', http_options=mldev_http_options
)
assert not mldev_client.models._api_client._httpx_client.trust_env
assert not mldev_client.models._api_client._async_httpx_client.trust_env
# Expect aiohttp to be disabled when httpx clients are provided, regardless of
# whether aiohttp is installed.
assert not mldev_client.models._api_client._use_aiohttp()
vertexai_client = Client(
vertexai=True,
project='fake_project_id',
location='fake-location',
http_options=vertexai_http_options,
)
assert not vertexai_client.models._api_client._httpx_client.trust_env
assert not vertexai_client.models._api_client._async_httpx_client.trust_env
# Expect aiohttp to be disabled when httpx clients are provided, regardless of
# whether aiohttp is installed.
assert not mldev_client.models._api_client._use_aiohttp()
# Aiohttp
@requires_aiohttp
@pytest.mark.asyncio
@pytest.mark.skipif(
AIOHTTP_NOT_INSTALLED, reason='aiohttp is not installed, skipping test.'
)
async def test_constructor_with_aiohttp_clients():
api_client.has_aiohttp = True
mldev_http_options = {
'aiohttp_client': aiohttp.ClientSession(trust_env=False),
}
vertexai_http_options = {
'aiohttp_client': aiohttp.ClientSession(trust_env=False),
}
mldev_client = Client(
api_key='google_api_key', http_options=mldev_http_options
)
assert not mldev_client.models._api_client._aiohttp_session.trust_env
vertexai_client = Client(
vertexai=True,
project='fake_project_id',
location='fake-location',
http_options=vertexai_http_options,
)
assert not vertexai_client.models._api_client._aiohttp_session.trust_env

View File

@@ -0,0 +1,182 @@
# 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.
#
"""Tests for client behavior when issuing requests."""
import pytest
from ... import _api_client
from ... import types
def test_patch_http_options_with_copies_all_fields():
patch_options = types.HttpOptions(
base_url='https://fake-url.com/',
api_version='v1',
headers={'X-Custom-Header': 'custom_value'},
timeout=10000,
client_args={'http2': True},
async_client_args={'http1': True},
extra_body={'key': 'value'},
retry_options=types.HttpRetryOptions(attempts=10),
base_url_resource_scope=types.ResourceScope.COLLECTION,
)
options = types.HttpOptions()
patched = _api_client.patch_http_options(options, patch_options)
http_options_keys = types.HttpOptions.model_fields.keys()
for key in http_options_keys:
assert hasattr(patched, key)
if key not in [
'httpx_client',
'httpx_async_client',
'aiohttp_client',
]:
assert getattr(patched, key) is not None
assert patched.base_url == 'https://fake-url.com/'
assert patched.api_version == 'v1'
assert patched.headers['X-Custom-Header'] == 'custom_value'
assert patched.timeout == 10000
assert patched.retry_options.attempts == 10
assert patched.client_args['http2']
assert patched.async_client_args['http1']
def test_patch_http_options_merges_headers():
original_options = types.HttpOptions(
headers={
'X-Custom-Header': 'different_value',
'X-different-header': 'different_value',
}
)
patch_options = types.HttpOptions(
base_url='https://fake-url.com/',
api_version='v1',
headers={'X-Custom-Header': 'custom_value'},
timeout=10000,
)
patched = _api_client.patch_http_options(original_options, patch_options)
# If the header is present in both the original and patch options, the patch
# options value should be used
assert patched.headers['X-Custom-Header'] == 'custom_value'
assert patched.headers['X-different-header'] == 'different_value'
assert patched.base_url == 'https://fake-url.com/'
assert patched.api_version == 'v1'
assert patched.timeout == 10000
def test_patch_http_options_appends_version_headers():
original_options = types.HttpOptions(
headers={
'X-Custom-Header': 'different_value',
'X-different-header': 'different_value',
}
)
patch_options = types.HttpOptions(
base_url='https://fake-url.com/',
api_version='v1',
headers={'X-Custom-Header': 'custom_value'},
timeout=10000,
)
patched = _api_client.patch_http_options(original_options, patch_options)
assert 'user-agent' in patched.headers
assert 'x-goog-api-client' in patched.headers
def test_setting_timeout_populates_server_timeout_header():
api_client = _api_client.BaseApiClient(
vertexai=False,
api_key='test_api_key',
http_options=types.HttpOptions(timeout=10000),
)
request = api_client._build_request(
http_method='POST',
path='sample/path',
request_dict={},
)
assert 'X-Server-Timeout' in request.headers
assert request.headers['X-Server-Timeout'] == '10'
def test_timeout_rounded_to_nearest_second():
api_client = _api_client.BaseApiClient(
vertexai=False,
api_key='test_api_key',
)
http_options = types.HttpOptions(timeout=7300)
request = api_client._build_request(
http_method='POST',
path='sample/path',
request_dict={},
http_options=http_options,
)
assert request.headers['X-Server-Timeout'] == '8'
def test_server_timeout_not_overwritten():
api_client = _api_client.BaseApiClient(
vertexai=False,
api_key='test_api_key',
)
http_options = types.HttpOptions(
headers={'X-Server-Timeout': '3'},
timeout=11000)
request = api_client._build_request(
http_method='POST',
path='sample/path',
request_dict={},
http_options=http_options,
)
assert request.headers['X-Server-Timeout'] == '3'
def test_server_timeout_not_set_by_default():
api_client = _api_client.BaseApiClient(
vertexai=False,
api_key='test_api_key',
)
request = api_client._build_request(
http_method='POST',
path='sample/path',
request_dict={},
)
assert not 'X-Server-Timeout' in request.headers
def test_resource_scope_without_base_url_raises_error():
with pytest.raises(ValueError):
_api_client.BaseApiClient(
vertexai=True,
http_options=types.HttpOptions(
base_url_resource_scope=types.ResourceScope.COLLECTION,
),
)
def test_base_url_resource_scope_not_set_by_default():
api_client = _api_client.BaseApiClient(
vertexai=True,
http_options=types.HttpOptions(
base_url='https://fake-url.com/',
),
)
assert api_client._http_options.base_url_resource_scope is None
def test_retry_options_not_set_by_default():
options = types.HttpOptions()
assert options.retry_options is None

View File

@@ -0,0 +1,168 @@
# 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.
#
"""Tests replay client correctly compares requests ignoring key casing."""
from ... import _replay_api_client
from ... import types
def test_equal_objects_with_same_casing_returns_true():
obj1 = [{
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
'generationConfig': {'topK': 2.0, 'maxOutputTokens': 3},
}]
obj2 = [{
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
'generationConfig': {'topK': 2.0, 'maxOutputTokens': 3},
}]
assert _replay_api_client._equals_ignore_key_case(obj1, obj2)
def test_equal_objects_with_different_casing_returns_true():
obj1 = [{
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
'generationConfig': {'topK': 2.0, 'maxOutputTokens': 3},
}]
obj2 = [{
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
'generation_config': {'topK': 2.0, 'max_output_tokens': 3},
}]
assert _replay_api_client._equals_ignore_key_case(obj1, obj2)
def test_equal_objects_with_different_nested_casing_returns_true():
obj1 = [{
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
'generationConfig': {'topK': 2.0, 'maxOutputTokens': 3},
}]
obj2 = [{
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
'generationConfig': {'top_k': 2.0, 'maxOutputTokens': 3},
}]
assert _replay_api_client._equals_ignore_key_case(obj1, obj2)
def test_equal_int_float_values_returns_true():
obj1 = [{
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
'generationConfig': {'topK': 2.0, 'maxOutputTokens': 3},
}]
obj2 = [{
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
'generation_config': {'topK': 2, 'max_output_tokens': 3},
}]
assert _replay_api_client._equals_ignore_key_case(obj1, obj2)
def test_equal_enum_string_values_returns_true():
obj1 = [{
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
'generationConfig': {
'topK': 2.0,
'maxOutputTokens': 3,
'type': types.Type.STRING,
},
}]
obj2 = [{
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
'generation_config': {
'topK': 2,
'max_output_tokens': 3,
'type': 'STRING',
},
}]
assert _replay_api_client._equals_ignore_key_case(obj1, obj2)
def test_equal_enum_values_returns_true():
obj1 = [{
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
'generationConfig': {
'topK': 2.0,
'maxOutputTokens': 3,
'type': types.Type.STRING,
},
}]
obj2 = [{
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
'generation_config': {
'topK': 2,
'max_output_tokens': 3,
'type': types.Type.STRING,
},
}]
assert _replay_api_client._equals_ignore_key_case(obj1, obj2)
def test_different_number_value_returns_false():
obj1 = [{
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
'generationConfig': {'topK': 3, 'maxOutputTokens': 3},
}]
obj2 = [{
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
'generation_config': {'topK': 2.0, 'max_output_tokens': 3},
}]
assert not _replay_api_client._equals_ignore_key_case(
obj1, obj2
)
def test_different_string_value_returns_false():
obj1 = [{
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
'generationConfig': {'topK': 3, 'maxOutputTokens': 3},
}]
obj2 = [{
'contents': [
{'parts': [{'text': 'What is your name?'}], 'role': 'model'}
],
'generation_config': {'topK': 3, 'max_output_tokens': 3},
}]
assert not _replay_api_client._equals_ignore_key_case(
obj1, obj2
)
def test_different_enum_values_returns_false():
obj1 = [{
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
'generationConfig': {
'topK': 2.0,
'maxOutputTokens': 3,
'type': types.Type.STRING,
},
}]
obj2 = [{
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
'generation_config': {
'topK': 2,
'max_output_tokens': 3,
'type': types.Type.OBJECT,
},
}]
assert not _replay_api_client._equals_ignore_key_case(
obj1, obj2
)

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,197 @@
# 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.
#
"""Tests for error handling in the client api client layer for uploads."""
import io
import json
from unittest import mock
import pytest
import httpx
from ... import _api_client as api_client
from ... import errors
from ... import types
try:
import aiohttp
AIOHTTP_NOT_INSTALLED = False
except ImportError:
AIOHTTP_NOT_INSTALLED = True
def _httpx_response(code: int, headers: dict = None, content: bytes = b''):
headers = headers or {}
headers['status-code'] = str(code)
return httpx.Response(
status_code=code,
headers=headers,
content=content,
)
@pytest.fixture
def client():
return api_client.BaseApiClient(vertexai=False, api_key='test_api_key')
def test_upload_url_rewrite(client: api_client.BaseApiClient):
mock_httpx_client = mock.MagicMock(spec=httpx.Client)
mock_httpx_client.request.side_effect = [
_httpx_response(
200, headers={"X-Goog-Upload-Status": "final"}
), # Upload request succeeding
]
client._httpx_client = mock_httpx_client
http_options = types.HttpOptions(base_url="https://my-proxy.company.com")
with io.BytesIO(b"test") as f:
client._upload_fd(
f,
"https://generativelanguage.googleapis.com/upload/v1beta/files?uploadType=resumable",
4,
http_options=http_options,
)
assert mock_httpx_client.request.call_count == 1
call_args = mock_httpx_client.request.call_args[1]
assert (
call_args["url"]
== "https://my-proxy.company.com/upload/v1beta/files?uploadType=resumable"
)
def test_upload_fd_error(client: api_client.BaseApiClient):
error_content = json.dumps({
"error": {
"code": 400,
"message": "Unsupported MIME type: bad/mime_type",
"status": "INVALID_ARGUMENT"
}
}).encode('utf-8')
mock_httpx_client = mock.MagicMock(spec=httpx.Client)
mock_httpx_client.request.side_effect = [
_httpx_response(200, headers={"X-Goog-Upload-URL": "http://fake/upload"}), # Initial request to get upload URL
_httpx_response(400, content=error_content, headers={"X-Goog-Upload-Status": "final"}), # Upload request failing
]
client._httpx_client = mock_httpx_client
with pytest.raises(errors.APIError, match='Unsupported MIME type: bad/mime_type'), io.BytesIO(b"test") as f:
client._upload_fd(
f, # Pass file object directly
"http://fake/upload",
4,
http_options=types.HttpOptions(),
)
assert mock_httpx_client.request.call_count == 2
@pytest.mark.asyncio
async def test_async_upload_url_rewrite_httpx(client: api_client.BaseApiClient):
mock_async_httpx_client = mock.MagicMock(spec=httpx.AsyncClient)
mock_async_httpx_client.request = mock.AsyncMock(
side_effect=[
_httpx_response(
200, headers={"X-Goog-Upload-Status": "final"}
), # Upload request
]
)
client._async_httpx_client = mock_async_httpx_client
http_options = types.HttpOptions(base_url="https://my-proxy.company.com")
with mock.patch.object(
client, "_use_aiohttp", return_value=False
), io.BytesIO(b"test") as f:
await client._async_upload_fd(
f,
"https://generativelanguage.googleapis.com/upload/v1beta/files?uploadType=resumable",
4,
http_options=http_options,
)
assert mock_async_httpx_client.request.call_count == 1
call_args = mock_async_httpx_client.request.call_args[1]
assert (
call_args["url"]
== "https://my-proxy.company.com/upload/v1beta/files?uploadType=resumable"
)
@pytest.mark.asyncio
async def test_async_upload_fd_error_httpx(client: api_client.BaseApiClient):
error_content = json.dumps({
"error": {
"code": 400,
"message": "Unsupported MIME type: bad/mime_type",
"status": "INVALID_ARGUMENT"
}
}).encode('utf-8')
# Mock for httpx.AsyncClient
mock_async_httpx_client = mock.MagicMock(spec=httpx.AsyncClient)
mock_async_httpx_client.request = mock.AsyncMock(side_effect=[
_httpx_response(200, headers={"X-Goog-Upload-URL": "http://fake/upload"}), # Initial request
_httpx_response(400, content=error_content, headers={"X-Goog-Upload-Status": "final"}), # Upload request
])
client._async_httpx_client = mock_async_httpx_client
# Patch the _use_aiohttp method to control which client is used
with mock.patch.object(client, '_use_aiohttp', return_value=False), \
pytest.raises(errors.APIError, match='Unsupported MIME type: bad/mime_type'), \
io.BytesIO(b"test") as f:
await client._async_upload_fd(
f, # Pass file object directly
"http://fake/upload",
4,
http_options=types.HttpOptions(),
)
assert mock_async_httpx_client.request.call_count == 2
@pytest.mark.asyncio
@pytest.mark.skipif(
AIOHTTP_NOT_INSTALLED, reason="aiohttp is not installed, skipping test."
)
async def test_async_upload_fd_error_aiohttp(client: api_client.BaseApiClient):
error_content = json.dumps({
"error": {
"code": 400,
"message": "Unsupported MIME type: bad/mime_type",
"status": "INVALID_ARGUMENT"
}
}).encode('utf-8')
# Mock for aiohttp.ClientSession
mock_aiohttp_session = mock.MagicMock()
mock_aiohttp_session.request = mock.AsyncMock(side_effect=[
_httpx_response(200, headers={"X-Goog-Upload-URL": "http://fake/upload"}), # Initial request
_httpx_response(400, content=error_content, headers={"X-Goog-Upload-Status": "final"}), # Upload request
])
with mock.patch.object(client, '_use_aiohttp', return_value=True), \
mock.patch.object(client, '_get_aiohttp_session', return_value=mock_aiohttp_session), \
pytest.raises(errors.APIError, match='Unsupported MIME type: bad/mime_type'), \
io.BytesIO(b"test") as f:
await client._async_upload_fd(
f, # Pass file object directly
"http://fake/upload",
4,
http_options=types.HttpOptions(),
)
assert mock_aiohttp_session.request.call_count == 2

View File

@@ -0,0 +1,17 @@
# 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.
#
"""Tests for the Google GenAI SDK's _common module."""

View File

@@ -0,0 +1,954 @@
# 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.
#
"""Tests tools in the _common module."""
from enum import Enum
import inspect
import logging
import textwrap
import typing
from typing import List, Optional
import warnings
import pydantic
import pytest
from ... import _common
from ... import types
from ... import errors
def test_warn_once():
@_common.experimental_warning('Warning!')
def func():
pass
with warnings.catch_warnings(record=True) as w:
func()
func()
assert len(w) == 1
assert w[0].category == errors.ExperimentalWarning
def test_warn_at_call_line():
@_common.experimental_warning('Warning!')
def func():
pass
with warnings.catch_warnings(record=True) as captured_warnings:
call_line = inspect.currentframe().f_lineno + 1
func()
assert captured_warnings[0].lineno == call_line
def test_is_struct_type():
assert _common._is_struct_type(list[dict[str, typing.Any]])
assert _common._is_struct_type(typing.List[typing.Dict[str, typing.Any]])
assert not _common._is_struct_type(list[dict[str, int]])
assert not _common._is_struct_type(list[dict[int, typing.Any]])
assert not _common._is_struct_type(list[str])
assert not _common._is_struct_type(dict[str, typing.Any])
assert not _common._is_struct_type(typing.List[typing.Dict[str, int]])
assert not _common._is_struct_type(typing.List[typing.Dict[int, typing.Any]])
assert not _common._is_struct_type(typing.List[str])
assert not _common._is_struct_type(typing.Dict[str, typing.Any])
@pytest.mark.parametrize(
"test_id, initial_target, update_dict, expected_target",
[
(
"simple_update",
{"a": 1, "b": 2},
{"b": 3, "c": 4},
{"a": 1, "b": 3, "c": 4},
),
(
"nested_update",
{"a": 1, "b": {"x": 10, "y": 20}},
{"b": {"y": 30, "z": 40}, "c": 3},
{"a": 1, "b": {"x": 10, "y": 30, "z": 40}, "c": 3},
),
(
"add_new_nested_dict",
{"a": 1},
{"b": {"x": 10, "y": 20}},
{"a": 1, "b": {"x": 10, "y": 20}},
),
(
"empty_target",
{},
{"a": 1, "b": {"x": 10}},
{"a": 1, "b": {"x": 10}},
),
(
"empty_update",
{"a": 1, "b": {"x": 10}},
{},
{"a": 1, "b": {"x": 10}},
),
(
"overwrite_non_dict_with_dict",
{"a": 1, "b": 2},
{"b": {"x": 10}},
{"a": 1, "b": {"x": 10}},
),
(
"overwrite_dict_with_non_dict",
{"a": 1, "b": {"x": 10}},
{"b": 2},
{"a": 1, "b": 2},
),
(
"deeper_nesting",
{"a": {"b": {"c": 1, "d": 2}, "e": 3}},
{"a": {"b": {"d": 4, "f": 5}, "g": 6}, "h": 7},
{"a": {"b": {"c": 1, "d": 4, "f": 5}, "e": 3, "g": 6}, "h": 7},
),
(
"different_value_types",
{"key1": "string_val", "key2": {"nested_int": 100}},
{"key1": 123, "key2": {"nested_list": [1, 2, 3]}, "key3": True},
{
"key1": 123,
"key2": {"nested_int": 100, "nested_list": [1, 2, 3]},
"key3": True,
},
),
(
"update_with_empty_nested_dict", # Existing nested dict in target should not be cleared
{"a": {"b": 1}},
{"a": {}},
{"a": {"b": 1}},
),
(
"target_with_empty_nested_dict",
{"a": {}},
{"a": {"b": 1}},
{"a": {"b": 1}},
),
(
"key_case_alignment_check",
{"first_name": "John", "contact_info": {"email_address": "john@example.com"}},
{"firstName": "Jane", "contact_info": {"email_address": "jane@example.com", "phone_number": "123"}},
{"first_name": "Jane", "contact_info": {"email_address": "jane@example.com", "phone_number": "123"}},
)
],
)
def test_recursive_dict_update(
test_id: str, initial_target: dict, update_dict: dict, expected_target: dict
):
_common.recursive_dict_update(initial_target, update_dict)
assert initial_target == expected_target
@pytest.mark.parametrize(
"test_id, initial_target, update_dict, expected_target, expect_warning, expected_log_message_part",
[
(
"type_match_int",
{"a": 1},
{"a": 2},
{"a": 2},
False,
"",
),
(
"type_match_dict",
{"a": {"b": 1}},
{"a": {"b": 2}},
{"a": {"b": 2}},
False,
"",
),
(
"type_mismatch_int_to_str",
{"a": 1},
{"a": "hello"},
{"a": "hello"},
True,
"Type mismatch for key 'a'. Existing type: <class 'int'>, new type: <class 'str'>. Overwriting.",
),
(
"type_mismatch_dict_to_int",
{"a": {"b": 1}},
{"a": 100},
{"a": 100},
True,
"Type mismatch for key 'a'. Existing type: <class 'dict'>, new type: <class 'int'>. Overwriting.",
),
(
"type_mismatch_int_to_dict",
{"a": 100},
{"a": {"b": 1}},
{"a": {"b": 1}},
True,
"Type mismatch for key 'a'. Existing type: <class 'int'>, new type: <class 'dict'>. Overwriting.",
),
("add_new_key", {"a": 1}, {"b": "new"}, {"a": 1, "b": "new"}, False, ""),
],
)
def test_recursive_dict_update_type_warnings(test_id, initial_target, update_dict, expected_target, expect_warning, expected_log_message_part, caplog):
_common.recursive_dict_update(initial_target, update_dict)
assert initial_target == expected_target
if expect_warning:
assert len(caplog.records) == 1
assert caplog.records[0].levelname == "WARNING"
assert expected_log_message_part in caplog.records[0].message
else:
for record in caplog.records:
if record.levelname == "WARNING" and expected_log_message_part in record.message:
pytest.fail(f"Unexpected warning logged for {test_id}: {record.message}")
@pytest.mark.parametrize(
"test_id, target_dict, update_dict, expected_aligned_dict",
[
(
"simple_snake_to_camel",
{"first_name": "John", "last_name": "Doe"},
{"firstName": "Jane", "lastName": "Doe"},
{"first_name": "Jane", "last_name": "Doe"},
),
(
"simple_camel_to_snake",
{"firstName": "John", "lastName": "Doe"},
{"first_name": "Jane", "last_name": "Doe"},
{"firstName": "Jane", "lastName": "Doe"},
),
(
"nested_dict_alignment",
{"user_info": {"contact_details": {"email_address": ""}}},
{"userInfo": {"contactDetails": {"emailAddress": "test@example.com"}}},
{"user_info": {"contact_details": {"email_address": "test@example.com"}}},
),
(
"list_of_dicts_alignment",
{"users_list": [{"user_id": 0, "user_name": ""}]},
{"usersList": [{"userId": 1, "userName": "Alice"}]},
{"users_list": [{"userId": 1, "userName": "Alice"}]},
),
(
"list_of_dicts_alignment_mixed_case_in_update",
{"users_list": [{"user_id": 0, "user_name": ""}]},
{"usersList": [{"user_id": 1, "UserName": "Alice"}]},
{"users_list": [{"user_id": 1, "UserName": "Alice"}]},
),
(
"list_of_dicts_different_lengths_update_longer",
{"items_data": [{"item_id": 0}]},
{"itemsData": [{"itemId": 1}, {"item_id": 2, "itemName": "Extra"}]},
{"items_data": [{"itemId": 1}, {"item_id": 2, "itemName": "Extra"}]},
),
(
"list_of_dicts_different_lengths_target_longer",
{"items_data": [{"item_id": 0, "item_name": ""}, {"item_id": 1}]},
{"itemsData": [{"itemId": 10}]},
{"items_data": [{"itemId": 10}]},
),
(
"no_matching_keys_preserves_update_case",
{"key_one": 1},
{"KEY_TWO": 2, "keyThree": 3},
{"KEY_TWO": 2, "keyThree": 3},
),
(
"mixed_match_and_no_match",
{"first_name": "John", "age_years": 30},
{"firstName": "Jane", "AGE_YEARS": 28, "occupation_title": "Engineer"},
{"first_name": "Jane", "age_years": 28, "occupation_title": "Engineer"},
),
(
"empty_target_dict",
{},
{"new_key": "new_value", "anotherKey": "anotherValue"},
{"new_key": "new_value", "anotherKey": "anotherValue"},
),
(
"empty_update_dict",
{"existing_key": "value"},
{},
{},
),
(
"target_has_non_dict_value_for_nested_key",
{"config_settings": 123},
{"configSettings": {"themeName": "dark"}},
{"config_settings": {"themeName": "dark"}}, # Overwrites as per recursive_dict_update logic
),
(
"update_has_non_dict_value_for_nested_key",
{"config_settings": {"theme_name": "light"}},
{"configSettings": "dark_theme_string"},
{"config_settings": "dark_theme_string"}, # Overwrites
),
(
"deeply_nested_with_lists",
{"level_one": {"list_items": [{"item_name": "", "item_value": 0}]}},
{"levelOne": {"listItems": [{"itemName": "Test", "itemValue": 100}, {"itemName": "Test2", "itemValue": 200}]}},
{"level_one": {"list_items": [{"itemName": "Test", "itemValue": 100}, {"itemName": "Test2", "itemValue": 200}]}},
),
],
)
def test_align_key_case(
test_id: str, target_dict: dict, update_dict: dict, expected_aligned_dict: dict
):
aligned_dict = _common.align_key_case(target_dict, update_dict)
assert aligned_dict == expected_aligned_dict, f"Test failed for: {test_id}"
class SimpleModel(_common.BaseModel):
name: str
value: int
is_active: bool = True
none_field: Optional[str] = None
class Chain(_common.BaseModel):
id: int
child: Optional["Chain"] = None
Chain.model_rebuild()
class Tree(_common.BaseModel):
id: int
children: List["Tree"] = pydantic.Field(default_factory=list)
Tree.model_rebuild()
class ReprFalseModel(_common.BaseModel):
visible: str
hidden: str = pydantic.Field("secret", repr=False)
class NonPydantic:
def __repr__(self):
return "NonPydantic(\n attr='value'\n)"
class MyEnum(Enum):
ONE = 1
TWO = 2
class EmptyModel(_common.BaseModel):
pass
def test_repr_simple_model_defaults_and_no_none():
obj = SimpleModel(name="Test Name", value=123)
expected = textwrap.dedent("""
SimpleModel(
is_active=True,
name='Test Name',
value=123
)
""").strip()
assert repr(obj) == expected
def test_repr_empty_model():
obj = EmptyModel()
expected = "EmptyModel()"
assert repr(obj) == expected
def test_repr_nested_model():
obj = Chain(id=1, child=Chain(id=2))
expected = textwrap.dedent("""
Chain(
child=Chain(
id=2
),
id=1
)
""").strip()
assert repr(obj) == expected
def test_repr_circular_model():
obj1 = Chain(id=1)
obj2 = Chain(id=2)
obj1.child = obj2
obj2.child = obj1 # Circular reference
expected = textwrap.dedent("""
Chain(
child=Chain(
child=<... Circular reference ...>,
id=2
),
id=1
)
""").strip()
assert repr(obj1) == expected
def test_repr_circular_list():
my_list = [1, 2]
my_list.append(my_list)
expected = textwrap.dedent("""
[
1,
2,
<... Circular reference ...>,
]
""").strip()
assert _common._pretty_repr(my_list) == expected
def test_repr_circular_dict():
my_dict = {"a": 1}
my_dict["self"] = my_dict
expected = textwrap.dedent("""
{
'a': 1,
'self': <... Circular reference ...>
}
""").strip()
assert _common._pretty_repr(my_dict) == expected
def test_repr_max_items():
lst = list(range(10))
dct = {i: i for i in range(10)}
st = set(range(10))
tpl = tuple(range(10))
assert (
"<... 5 more items ...>" in
_common._pretty_repr(lst, max_items=5)
)
assert (
"<dict len=10>" in _common._pretty_repr(dct, max_items=5))
assert (
"<... 5 more items ...>" in _common._pretty_repr(st, max_items=5)
)
assert (
"<... 5 more items ...>" in _common._pretty_repr(tpl, max_items=5)
)
def test_repr_max_len_bytes():
b_data = b"a" * 100
assert len(_common._pretty_repr(b_data, max_len=90)) == 90 + 3
assert repr(b_data) == _common._pretty_repr(b_data, max_len=200)
def test_repr_max_depth_dict():
nested = {'a': {'a': {'a': {'a': 'a', 'b': 'b'}}}}
assert "{<... 2 items at Max depth ...>}" in _common._pretty_repr(nested, depth=3)
def test_repr_max_depth_list():
nested = [[[["d", "e", "e", "p"]]]]
assert "[<... 4 items at Max depth ...>]" in _common._pretty_repr(nested, depth=3)
def test_repr_collections():
obj = {
"set": {3, 1, 2},
"tuple": (4, 5, 6),
"dict": {"b": 2, "a": 1},
"list": [7, 8, 9],
}
expected = textwrap.dedent("""
{
'dict': {
'a': 1,
'b': 2
},
'list': [
7,
8,
9,
],
'set': {
1,
2,
3,
},
'tuple': (
4,
5,
6,
)
}
""").strip()
assert _common._pretty_repr(obj) == expected
def test_tuple_collections():
obj = {
"tuple0": (),
"tuple1": (1,),
"tuple2": (1, 2),
}
expected = textwrap.dedent("""
{
'tuple0': (),
'tuple1': (
1,
),
'tuple2': (
1,
2,
)
}
""").strip()
assert _common._pretty_repr(obj) == expected
def test_repr_empty_collections():
assert _common._pretty_repr([]) == "[]"
assert _common._pretty_repr({}) == "{}"
assert (
_common._pretty_repr(set()) == "set()"
)
assert _common._pretty_repr(tuple()) == "()"
assert (
_common._pretty_repr({"empty_set": set()}) ==
textwrap.dedent("""
{
'empty_set': set()
}
""").strip()
)
def test_repr_strings():
s1 = "line one"
exp1 = "'line one'"
assert _common._pretty_repr(s1) == exp1
s2 = 'line one\nline two with """ inside'
exp2 = '"""line one\nline two with \\"\\"\\" inside"""'
assert _common._pretty_repr(s2) == exp2
s3 = 'A string with """ inside'
exp3 = '\'A string with """ inside\''
assert _common._pretty_repr(s3) == exp3
def test_repr_repr_false():
obj = ReprFalseModel(visible="show", hidden="hide")
result = repr(obj)
assert "visible='show'" in result
assert "hidden" not in result
expected = textwrap.dedent("""
ReprFalseModel(
visible='show'
)
""").strip()
assert result == expected
def test_repr_none_fields():
obj = SimpleModel(name="Only Name", value=0, none_field=None)
result = repr(obj)
assert "none_field" not in result
expected = textwrap.dedent("""
SimpleModel(
is_active=True,
name='Only Name',
value=0
)
""").strip()
assert result == expected
def test_repr_other_types():
np = NonPydantic()
en = MyEnum.TWO
obj = {"np": np, "en": en}
expected = textwrap.dedent("""
{
'en': <MyEnum.TWO: 2>,
'np': NonPydantic(
attr='value'
)
}
""").strip()
assert _common._pretty_repr(obj) == expected
def test_repr_indent_delta():
obj = SimpleModel(name="Indent Test", value=1)
expected = textwrap.dedent("""
SimpleModel(
is_active=True,
name='Indent Test',
value=1
)
""").strip()
assert _common._pretty_repr(obj, indent_delta=4) == expected
def test_repr_complex_object():
obj = types.GenerateContentResponse(
automatic_function_calling_history=[],
candidates=[
types.Candidate(
content=types.Content(
parts=[
types.Part(
text="""There isn't a single "best" LLM, as the ideal choice highly depends on your specific needs, use case, budget, and priorities. The field is evolving incredibly fast, with new models and improvements being released constantly.
However, we can talk about the **leading contenders** and what they are generally known for:..."""
)
],
role="model"
),
finish_reason=types.FinishReason.STOP,
index=0
)
],
model_version='models/gemini-2.5-flash-preview-05-20',
usage_metadata=types.GenerateContentResponseUsageMetadata(
candidates_token_count=1086,
prompt_token_count=7,
prompt_tokens_details=[
types.ModalityTokenCount(
modality=types.MediaModality.TEXT,
token_count=7
)
],
thoughts_token_count=860,
total_token_count=1953
)
)
expected = textwrap.dedent("""
GenerateContentResponse(
automatic_function_calling_history=[],
candidates=[
Candidate(
content=Content(
parts=[
Part(
text=\"\"\"There isn't a single "best" LLM, as the ideal choice highly depends on your specific needs, use case, budget, and priorities. The field is evolving incredibly fast, with new models and improvements being released constantly.
However, we can talk about the **leading contenders** and what they are generally known for:...\"\"\"
),
],
role='model'
),
finish_reason=<FinishReason.STOP: 'STOP'>,
index=0
),
],
model_version='models/gemini-2.5-flash-preview-05-20',
usage_metadata=GenerateContentResponseUsageMetadata(
candidates_token_count=1086,
prompt_token_count=7,
prompt_tokens_details=[
ModalityTokenCount(
modality=<MediaModality.TEXT: 'TEXT'>,
token_count=7
),
],
thoughts_token_count=860,
total_token_count=1953
)
)
""").strip()
assert repr(obj) == expected
def test_move_value_by_path():
"""Test move_value_by_path function with array wildcard notation."""
data = {
"requests": [
{
"request": {
"content": {
"parts": [
{
"text": "1"
}
]
}
},
"outputDimensionality": 64
},
{
"request": {
"content": {
"parts": [
{
"text": "2"
}
]
}
},
"outputDimensionality": 64
},
{
"request": {
"content": {
"parts": [
{
"text": "3"
}
]
}
},
"outputDimensionality": 64
}
]
}
paths = {'requests[].*': 'requests[].request.*'}
_common.move_value_by_path(data, paths)
expected = {
"requests": [
{
"request": {
"content": {
"parts": [
{
"text": "1"
}
]
},
"outputDimensionality": 64
}
},
{
"request": {
"content": {
"parts": [
{
"text": "2"
}
]
},
"outputDimensionality": 64
}
},
{
"request": {
"content": {
"parts": [
{
"text": "3"
}
]
},
"outputDimensionality": 64
}
}
]
}
assert data == expected
def test_check_field_type_mismatches_no_warning_for_correct_types(caplog):
"""Test that no warning is logged when types match."""
class ModelA(_common.BaseModel):
value: int
class TestModel(_common.BaseModel):
model_a: ModelA
# Should not warn - dict will be converted to ModelA by Pydantic
data = {"model_a": {"value": 123}}
with caplog.at_level(logging.WARNING, logger="google_genai._common"):
result = TestModel.model_validate(data)
assert result.model_a.value == 123
assert len(caplog.records) == 0
def test_check_field_type_mismatches_warns_on_pydantic_type_mismatch(caplog):
"""Test that warning is logged when Pydantic model types mismatch."""
class ModelA(_common.BaseModel):
value: int
class ModelB(_common.BaseModel):
value: str
class TestModel(_common.BaseModel):
model_field: ModelA
# Create an instance of ModelB (wrong type)
model_b_instance = ModelB(value="test")
# Pass the wrong Pydantic model instance
data = {"model_field": model_b_instance}
with caplog.at_level(logging.WARNING, logger="google_genai._common"):
TestModel._check_field_type_mismatches(data)
assert len(caplog.records) == 1
assert "Type mismatch in TestModel.model_field" in caplog.records[0].message
assert "expected ModelA, got ModelB" in caplog.records[0].message
def test_check_field_type_mismatches_no_warning_for_none_values(caplog):
"""Test that no warning is logged for None values."""
class ModelA(_common.BaseModel):
value: int
class TestModel(_common.BaseModel):
model_field: Optional[ModelA] = None
data = {"model_field": None}
with caplog.at_level(logging.WARNING, logger="google_genai._common"):
result = TestModel.model_validate(data)
assert result.model_field is None
assert len(caplog.records) == 0
def test_check_field_type_mismatches_no_warning_for_missing_fields(caplog):
"""Test that no warning is logged for missing fields."""
class ModelA(_common.BaseModel):
value: int
class TestModel(_common.BaseModel):
model_field: Optional[ModelA] = None
data = {}
with caplog.at_level(logging.WARNING, logger="google_genai._common"):
result = TestModel.model_validate(data)
assert result.model_field is None
assert len(caplog.records) == 0
def test_check_field_type_mismatches_no_warning_for_primitive_types(caplog):
"""Test that no warning is logged for primitive type mismatches."""
class TestModel(_common.BaseModel):
int_field: int
str_field: str
# Even though we're passing wrong primitive types, we should not warn
# (Pydantic will handle validation)
data = {"int_field": "123", "str_field": "test"}
with caplog.at_level(logging.WARNING, logger="google_genai._common"):
# This will succeed because Pydantic can coerce "123" to int
result = TestModel.model_validate(data)
assert result.int_field == 123
assert result.str_field == "test"
assert len(caplog.records) == 0
def test_check_field_type_mismatches_handles_optional_unwrapping(caplog):
"""Test that Optional types are properly unwrapped before checking."""
class ModelA(_common.BaseModel):
value: int
class ModelB(_common.BaseModel):
value: str
class TestModel(_common.BaseModel):
model_field: Optional[ModelA] = None
# Pass wrong Pydantic model type
model_b_instance = ModelB(value="test")
data = {"model_field": model_b_instance}
with caplog.at_level(logging.WARNING, logger="google_genai._common"):
TestModel._check_field_type_mismatches(data)
assert len(caplog.records) == 1
assert "expected ModelA, got ModelB" in caplog.records[0].message
def test_check_field_type_mismatches_no_warning_for_correct_pydantic_instance(caplog):
"""Test that no warning is logged when correct Pydantic instance is provided."""
class ModelA(_common.BaseModel):
value: int
class TestModel(_common.BaseModel):
model_field: ModelA
# Pass correct Pydantic model instance
model_a_instance = ModelA(value=42)
data = {"model_field": model_a_instance}
with caplog.at_level(logging.WARNING, logger="google_genai._common"):
result = TestModel.model_validate(data)
assert result.model_field.value == 42
assert len(caplog.records) == 0
def test_check_field_type_mismatches_with_multiple_fields(caplog):
"""Test checking multiple fields with mixed scenarios."""
class ModelA(_common.BaseModel):
value: int
class ModelB(_common.BaseModel):
value: str
class TestModel(_common.BaseModel):
field_a: ModelA
field_b: Optional[ModelA] = None
field_c: str
model_b_instance = ModelB(value="wrong")
data = {
"field_a": model_b_instance, # Wrong type - should warn
"field_b": None, # None - should not warn
"field_c": "test", # Primitive - should not warn
}
with caplog.at_level(logging.WARNING, logger="google_genai._common"):
TestModel._check_field_type_mismatches(data)
# Should only warn about field_a
assert len(caplog.records) == 1
assert "field_a" in caplog.records[0].message
assert "expected ModelA, got ModelB" in caplog.records[0].message
def test_check_field_type_mismatches_generic_type_no_error(caplog):
"""Test that validation doesn't crash on generic types like list[str]."""
class TestModel(_common.BaseModel):
tags: list[str]
data = {"tags": ["a", "b"]}
with caplog.at_level(logging.WARNING, logger="google_genai._common"):
TestModel.model_validate(data)
assert len(caplog.records) == 0

View File

@@ -0,0 +1,96 @@
import unittest
import pydantic
from ... import _common
class TestIsDuckTypeOf(unittest.TestCase):
class FakePydanticModel(pydantic.BaseModel):
field1: str
field2: int
field3: str
field4: str
field5: str
class FakePydanticModelWithLessFields(pydantic.BaseModel):
field1: str
field2: int
field3: str
field4: str
def test_is_duck_type_of_true_for_pydantic_object(self):
obj = self.FakePydanticModel(
field1="a",
field2=1,
field3="b",
field4="c",
field5="d",
)
self.assertTrue(_common.is_duck_type_of(obj, self.FakePydanticModel))
def test_is_duck_type_of_true_for_duck_typed_object(self):
class DuckTypedObject:
def __init__(self):
self.field1 = "a"
self.field2 = 1
self.field3 = "b"
self.field4 = "c"
self.field5 = "d"
obj = DuckTypedObject()
self.assertTrue(_common.is_duck_type_of(obj, self.FakePydanticModel))
def test_is_duck_type_of_false_for_different_many_fields(self):
class DifferentFieldsObject:
def __init__(self):
self.fielda = "a"
obj = DifferentFieldsObject()
self.assertFalse(
_common.is_duck_type_of(obj, self.FakePydanticModel)
)
def test_is_duck_type_of_false_for_missing_fields(self):
obj = self.FakePydanticModelWithLessFields(
field1="a", field2=1, field3="b", field4="c"
)
self.assertFalse(_common.is_duck_type_of(obj, self.FakePydanticModel))
def test_is_duck_type_of_false_for_dict(self):
obj = {"field1": "a", "field2": 1}
self.assertFalse(
_common.is_duck_type_of(obj, self.FakePydanticModel)
)
def test_is_duck_type_of_false_for_non_pydantic_class(self):
class NonPydanticModel:
pass
class SomeObject:
pass
obj = SomeObject()
self.assertFalse(_common.is_duck_type_of(obj, NonPydanticModel))
def test_is_duck_type_of_true_with_extra_fields(self):
class ExtraFieldsObject:
def __init__(self):
self.field1 = "a"
self.field2 = 1
self.field3 = "extra"
self.field4 = "extra"
self.field5 = "extra"
self.field6 = "extra"
obj = ExtraFieldsObject()
self.assertTrue(_common.is_duck_type_of(obj, self.FakePydanticModel))
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,162 @@
# 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.
#
"""Conftest for google.genai tests."""
import datetime
import os
from unittest import mock
import uuid
import pytest
from .. import _common
from .. import _replay_api_client
from .. import client as google_genai_client_module
def pytest_addoption(parser):
parser.addoption(
'--mode',
action='store',
default='auto',
help="""Replay mode.
One of:
* auto: Replay if replay files exist, otherwise record.
* record: Always call the API and record.
* replay: Always replay, fail if replay files do not exist.
* api: Always call the API and do not record.
* tap: Always replay, fail if replay files do not exist. Also sets default values for the API key and replay directory.
""",
)
parser.addoption(
'--private',
action='store_true',
default=False,
help='Run private tests.',
)
# Overridden via parameterized test.
@pytest.fixture
def use_vertex():
return False
# Overridden at the module level for each test file.
@pytest.fixture
def replays_prefix():
return 'test'
def _get_replay_id(use_vertex: bool, replays_prefix: str) -> str:
test_name_ending = os.environ.get('PYTEST_CURRENT_TEST').split('::')[-1]
test_name = (
test_name_ending.split(' ')[0].split('[')[0]
+ '.'
+ ('vertex' if use_vertex else 'mldev')
)
return '/'.join([replays_prefix, test_name])
@pytest.fixture
def client(use_vertex, replays_prefix, http_options, request):
mode = request.config.getoption('--mode')
if mode not in ['auto', 'record', 'replay', 'api', 'tap']:
raise ValueError('Invalid mode: ' + mode)
test_function_name = request.function.__name__
test_filename = os.path.splitext(os.path.basename(request.path))[0]
if test_function_name.startswith(test_filename):
raise ValueError(f"""
{test_function_name}:
Do not include the test filename in the test function name.
keep the test function name short.""")
for word in ['mldev', 'ml_dev', 'vertex']:
if word in test_function_name:
raise ValueError(f"""
{test_function_name}:
You should never ever include api types in the file name (mldev or vertex).
All tests should run against all apis.
Assert an exception if the test is not supported in an API.""")
replay_id = _get_replay_id(use_vertex, replays_prefix)
if mode == 'tap':
mode = 'replay'
# Set various environment variables to ensure that the test runs.
os.environ['GOOGLE_API_KEY'] = 'dummy-api-key'
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = os.path.join(
os.path.dirname(__file__),
'credentials.json',
)
os.environ['GOOGLE_CLOUD_PROJECT'] = 'project-id'
os.environ['GOOGLE_CLOUD_LOCATION'] = 'location'
# Set the replay directory to the root directory of the replays.
# This is needed to ensure that the replay files are found.
replays_root_directory = os.path.abspath(
os.path.join(
os.path.dirname(__file__),
'../../../../../google/cloud/aiplatform/sdk/genai/replays',
)
)
os.environ['GOOGLE_GENAI_REPLAYS_DIRECTORY'] = replays_root_directory
# Get private arg.
private = request.config.getoption('--private')
replay_client = _replay_api_client.ReplayApiClient(
mode=mode,
replay_id=replay_id,
vertexai=use_vertex,
http_options=http_options,
private=private,
)
with mock.patch.object(
google_genai_client_module.Client, '_get_api_client'
) as patch_method:
patch_method.return_value = replay_client
google_genai_client = google_genai_client_module.Client(vertexai=use_vertex)
# Yield the client so that cleanup can be completed at the end of the test.
yield google_genai_client
# Save the replay after the test if we were in recording mode.
google_genai_client._api_client.close()
@pytest.fixture
def mock_timestamped_unique_name():
with mock.patch.object(
_common,
'timestamped_unique_name',
return_value='20240101000000_bd656',
) as unique_name_mock:
yield unique_name_mock
@pytest.fixture
def image_jpeg():
import PIL.Image
image_path = os.path.join(os.path.dirname(__file__), 'data', 'google.jpg')
with PIL.Image.open(image_path) as img:
yield img
@pytest.fixture
def image_png():
import PIL.Image
image_path = os.path.join(os.path.dirname(__file__), 'data', 'google.png')
with PIL.Image.open(image_path) as img:
yield img

View File

@@ -0,0 +1,17 @@
# 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.
#
"""Tests for the Google GenAI SDK's documents module."""

View File

@@ -0,0 +1,51 @@
# 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.
#
"""Tests for documents.delete()."""
import pytest
from ... import types
from .. import pytest_helper
# All tests will be run for both Vertex and MLDev.
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name="test_delete",
parameters=types._DeleteDocumentParameters(
name=(
"fileSearchStores/fr3l0ri2so25-a3r1ump9x821/documents/asurveyofmodernistpoetrytxt-01zyb2rig5gu"
),
config=types.DeleteDocumentConfig(force=True),
),
exception_if_vertex="only supported in the Gemini Developer client",
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method="file_search_stores.documents.delete",
test_table=test_table,
)
@pytest.mark.asyncio
async def test_async_delete(client):
with pytest_helper.exception_if_vertex(client, ValueError):
await client.aio.file_search_stores.documents.delete(
name="fileSearchStores/fr3l0ri2so25-a3r1ump9x821/documents/asurveyofmodernistpoetrytxt-uvmqjtmkm1h2",
config=types.DeleteDocumentConfig(force=True),
)

View File

@@ -0,0 +1,85 @@
# 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.
#
"""Tests for file_search_stores.documents.get()."""
import pytest
from ... import types
from .. import pytest_helper
# A Document name known to exist in the test environment.
# Replace with a real one from your test setup.
_EXISTING_DOCUMENT_NAME = "fileSearchStores/fr3l0ri2so25-a3r1ump9x821/documents/asurveyofmodernistpoetrytxt-uvmqjtmkm1h2"
_NON_EXISTENT_DOCUMENT_NAME = (
"fileSearchStores/fr3l0ri2so25-a3r1ump9x821/documents/non-existent-document"
)
_INVALID_DOCUMENT_NAME = "fileSearchStores/fr3l0ri2so25-a3r1ump9x821/documents/_invalid_document_name"
_NOT_A_DOCUMENT_NAME = "genai-test-document"
# All tests will be run for both Vertex and MLDev.
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name="test_get_success",
parameters=types._GetDocumentParameters(name=_EXISTING_DOCUMENT_NAME),
exception_if_vertex="only supported in the Gemini Developer client",
),
pytest_helper.TestTableItem(
name="test_get_not_found",
parameters=types._GetDocumentParameters(
name=_NON_EXISTENT_DOCUMENT_NAME
),
exception_if_vertex="only supported in the Gemini Developer client",
exception_if_mldev="Documents does not exist",
),
pytest_helper.TestTableItem(
name="test_get_invalid_name",
parameters=types._GetDocumentParameters(name=_INVALID_DOCUMENT_NAME),
exception_if_vertex="only supported in the Gemini Developer client",
# Validation should catch this before the API call
exception_if_mldev="INVALID_ARGUMENT",
),
pytest_helper.TestTableItem(
name="test_get_not_a_document_name",
parameters=types._GetDocumentParameters(name=_NOT_A_DOCUMENT_NAME),
exception_if_vertex="only supported in the Gemini Developer client",
exception_if_mldev="Not Found",
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method="file_search_stores.documents.get",
test_table=test_table,
)
@pytest.mark.asyncio
async def test_async_get(client):
with pytest_helper.exception_if_vertex(client, ValueError):
# This test relies on the mocking/replays set up by pytest_helper
# For a success case:
try:
document = await client.aio.file_search_stores.documents.get(
name=_EXISTING_DOCUMENT_NAME
)
assert document is not None
assert document.name == _EXISTING_DOCUMENT_NAME
except Exception as e:
# Depending on the mock setup, errors might be raised directly
print(f"Async get failed: {e}")
raise

View File

@@ -0,0 +1,72 @@
# 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.
#
"""Tests for file_search_stores.documents.list()."""
import pytest
from ... import types
from .. import pytest_helper
_FILE_SEARCH_STORE_NAME = "fileSearchStores/gzn7kdl2wpxl-4z2yqvuxbcxw"
# All tests will be run for both Vertex and MLDev.
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name="test_list_default",
parameters=types._ListDocumentsParameters(
parent=_FILE_SEARCH_STORE_NAME
),
exception_if_vertex="only supported in the Gemini Developer client",
),
pytest_helper.TestTableItem(
name="test_list_with_page_size",
parameters=types._ListDocumentsParameters(
parent=_FILE_SEARCH_STORE_NAME,
config=types.ListDocumentsConfig(
page_size=2,
),
),
exception_if_vertex="only supported in the Gemini Developer client",
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method="file_search_stores.documents.list",
test_table=test_table,
http_options={
"base_url": "https://autopush-generativelanguage.sandbox.googleapis.com/",
}
)
def test_pager(client):
with pytest_helper.exception_if_vertex(client, ValueError):
# Iterate through the first page.
for document in client.file_search_stores.documents.list(
parent=_FILE_SEARCH_STORE_NAME, config={"page_size": 2}
):
assert isinstance(document, types.Document)
@pytest.mark.asyncio
async def test_async_pager(client):
with pytest_helper.exception_if_vertex(client, ValueError):
# Iterate through the first page.
async for document in await client.aio.file_search_stores.documents.list(
parent=_FILE_SEARCH_STORE_NAME, config={"page_size": 2}
):
assert isinstance(document, types.Document)

View File

@@ -0,0 +1 @@
"""Unit Tests for the error modules."""

View File

@@ -0,0 +1,454 @@
# 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.
#
"""Unit Tests for the APIError class.
End to end tests should be in models/test_generate_content.py.
"""
import pickle
import httpx
import pytest
from ... import errors
def test_constructor_code_none_error_in_json_code_in_error():
actual_error = errors.APIError(
None,
{
'error': {
'code': 400,
'message': 'error message',
'status': 'INVALID_ARGUMENT',
}
},
httpx.Response(status_code=400),
)
assert actual_error.code == 400
assert actual_error.message == 'error message'
assert actual_error.status == 'INVALID_ARGUMENT'
assert actual_error.details == {
'error': {
'code': 400,
'message': 'error message',
'status': 'INVALID_ARGUMENT',
}
}
def test_constructor_code_none_error_in_json_code_outside_error():
actual_error = errors.APIError(
None,
{
'code': 400,
'error': {
'code': 500,
'message': 'error message',
'status': 'INVALID_ARGUMENT',
},
},
httpx.Response(status_code=400),
)
assert actual_error.code == 400
assert actual_error.message == 'error message'
assert actual_error.status == 'INVALID_ARGUMENT'
assert actual_error.details == {
'code': 400,
'error': {
'code': 500,
'message': 'error message',
'status': 'INVALID_ARGUMENT',
},
}
def test_constructor_code_not_present():
actual_error = errors.APIError(
None,
{
'error': {
'message': 'error message',
'status': 'INVALID_ARGUMENT',
}
},
httpx.Response(status_code=400),
)
assert actual_error.code is None
assert actual_error.message == 'error message'
assert actual_error.status == 'INVALID_ARGUMENT'
assert actual_error.details == {
'error': {
'message': 'error message',
'status': 'INVALID_ARGUMENT',
}
}
def test_constructor_code_exist_error_in_json():
actual_error = errors.APIError(
400,
{
'error': {
'code': 400,
'message': 'error message',
'status': 'INVALID_ARGUMENT',
}
},
httpx.Response(status_code=400),
)
assert actual_error.code == 400
assert actual_error.message == 'error message'
assert actual_error.status == 'INVALID_ARGUMENT'
assert actual_error.details == {
'error': {
'code': 400,
'message': 'error message',
'status': 'INVALID_ARGUMENT',
}
}
def test_constructor_error_not_in_json():
actual_error = errors.APIError(
400,
{
'message': 'error message',
'status': 'INVALID_ARGUMENT',
'code': 400,
},
httpx.Response(status_code=400),
)
assert actual_error.code == 400
assert actual_error.message == 'error message'
assert actual_error.status == 'INVALID_ARGUMENT'
assert actual_error.details == {
'message': 'error message',
'status': 'INVALID_ARGUMENT',
'code': 400,
}
def test_constructor_error_in_json_status_outside_error():
actual_error = errors.APIError(
400,
{
'status': 'OUTER_INVALID_ARGUMENT_STATUS',
'error': {
'code': 400,
'message': 'error message',
'status': 'INNER_INVALID_ARGUMENT_STATUS',
},
},
httpx.Response(status_code=400),
)
assert actual_error.code == 400
assert actual_error.message == 'error message'
assert actual_error.status == 'OUTER_INVALID_ARGUMENT_STATUS'
assert actual_error.details == {
'status': 'OUTER_INVALID_ARGUMENT_STATUS',
'error': {
'code': 400,
'message': 'error message',
'status': 'INNER_INVALID_ARGUMENT_STATUS',
},
}
def test_constructor_status_not_present():
actual_error = errors.APIError(
400,
{
'error': {
'code': 400,
'message': 'error message',
}
},
httpx.Response(status_code=400),
)
assert actual_error.code == 400
assert actual_error.message == 'error message'
assert actual_error.status == None
assert actual_error.details == {
'error': {
'code': 400,
'message': 'error message',
}
}
def test_constructor_error_in_json_message_outside_error():
actual_error = errors.APIError(
400,
{
'message': 'OUTER_ERROR_MESSAGE',
'error': {
'code': 400,
'message': 'INNER_ERROR_MESSAGE',
'status': 'INVALID_ARGUMENT',
},
},
httpx.Response(status_code=400),
)
assert actual_error.code == 400
assert actual_error.message == 'OUTER_ERROR_MESSAGE'
assert actual_error.status == 'INVALID_ARGUMENT'
assert actual_error.details == {
'message': 'OUTER_ERROR_MESSAGE',
'error': {
'code': 400,
'message': 'INNER_ERROR_MESSAGE',
'status': 'INVALID_ARGUMENT',
},
}
def test_constructor_message_not_present():
actual_error = errors.APIError(
400,
{
'error': {
'code': 400,
'status': 'INVALID_ARGUMENT',
}
},
httpx.Response(status_code=400),
)
assert actual_error.code == 400
assert actual_error.message is None
assert actual_error.status == 'INVALID_ARGUMENT'
assert actual_error.details == {
'error': {
'code': 400,
'status': 'INVALID_ARGUMENT',
}
}
def test_constructor_with_websocket_connection_closed_error():
actual_error = errors.APIError(
1007,
'At most one response modality can be specified in the setup request.'
' To enable simultaneous transcription and audio output,',
None,
)
assert actual_error.code == 1007
assert (
actual_error.details
== 'At most one response modality can be specified in the setup request.'
' To enable simultaneous transcription and audio output,'
)
assert actual_error.status == None
assert actual_error.message == None
def test_raise_for_websocket_connection_closed_error():
try:
errors.APIError.raise_error(
1007,
'At most one response modality can be specified in the setup request.'
' To enable simultaneous transcription and audio output,',
None,
)
except errors.APIError as actual_error:
assert actual_error.code == 1007
assert (
actual_error.details
== 'At most one response modality can be specified in the setup'
' request.'
' To enable simultaneous transcription and audio output,'
)
assert actual_error.status == None
assert actual_error.message == None
def test_raise_for_response_code_exist_json_decoder_error():
class FakeResponse(httpx.Response):
def read(self) -> bytes:
self._content = b'{"data": {"key1": "value1", "key2"}'
return self._content
try:
errors.APIError.raise_for_response(
FakeResponse(
status_code=503,
extensions={'reason_phrase': b'Service Unavailable'},
)
)
except errors.ServerError as actual_error:
assert actual_error.code == 503
assert actual_error.message == '{"data": {"key1": "value1", "key2"}'
assert actual_error.status == 'Service Unavailable'
assert actual_error.details == {
'message': '{"data": {"key1": "value1", "key2"}',
'status': 'Service Unavailable',
}
def test_raise_for_response_client_error():
class FakeResponse(httpx.Response):
def read(self) -> bytes:
self._content = (
b'{"error": {"code": 400, "message": "error message", "status":'
b' "INVALID_ARGUMENT"}}'
)
return self._content
try:
errors.APIError.raise_for_response(FakeResponse(status_code=400))
except errors.ClientError as actual_error:
assert actual_error.code == 400
assert actual_error.message == 'error message'
assert actual_error.status == 'INVALID_ARGUMENT'
assert actual_error.details == {
'error': {
'code': 400,
'message': 'error message',
'status': 'INVALID_ARGUMENT',
}
}
def test_raise_for_response_server_error():
class FakeResponse(httpx.Response):
def read(self) -> bytes:
self._content = (
b'{"error": {"code": 500, "message": "error message", "status":'
b' "SERVER_INTERNAL ERROR"}}'
)
return self._content
try:
errors.APIError.raise_for_response(FakeResponse(status_code=500))
except errors.ServerError as actual_error:
assert actual_error.code == 500
assert actual_error.message == 'error message'
assert actual_error.status == 'SERVER_INTERNAL ERROR'
assert actual_error.details == {
'error': {
'code': 500,
'message': 'error message',
'status': 'SERVER_INTERNAL ERROR',
}
}
def test_api_error_is_picklable():
pickled_error = pickle.loads(pickle.dumps(errors.APIError(1, {})))
assert isinstance(pickled_error, errors.APIError)
@pytest.mark.asyncio
async def test_raise_for_async_response_client_error():
class FakeResponse(httpx.Response):
async def aread(self) -> bytes:
self._content = (
b'{"error": {"code": 400, "message": "error message", "status":'
b' "INVALID_ARGUMENT"}}'
)
return self._content
try:
await errors.APIError.raise_for_async_response(
FakeResponse(status_code=400)
)
except errors.ClientError as actual_error:
assert actual_error.code == 400
assert actual_error.message == 'error message'
assert actual_error.status == 'INVALID_ARGUMENT'
assert actual_error.details == {
'error': {
'code': 400,
'message': 'error message',
'status': 'INVALID_ARGUMENT',
}
}
@pytest.mark.asyncio
async def test_raise_for_async_response_server_error():
class FakeResponse(httpx.Response):
async def aread(self) -> bytes:
self._content = (
b'{"error": {"code": 500, "message": "error message", "status":'
b' "SERVER_INTERNAL ERROR"}}'
)
return self._content
try:
await errors.APIError.raise_for_async_response(
FakeResponse(status_code=500)
)
except errors.ServerError as actual_error:
assert actual_error.code == 500
assert actual_error.message == 'error message'
assert actual_error.status == 'SERVER_INTERNAL ERROR'
assert actual_error.details == {
'error': {
'code': 500,
'message': 'error message',
'status': 'SERVER_INTERNAL ERROR',
}
}
@pytest.mark.asyncio
async def test_raise_for_async_response_code_exist_json_decoder_error():
class FakeResponse(httpx.Response):
async def aread(self) -> bytes:
self._content = b'{"data": {"key1": "value1", "key2"}'
return self._content
try:
await errors.APIError.raise_for_async_response(
FakeResponse(
status_code=503,
extensions={'reason_phrase': b'Service Unavailable'},
)
)
except errors.ServerError as actual_error:
assert actual_error.code == 503
assert actual_error.message == '{"data": {"key1": "value1", "key2"}'
assert actual_error.status == 'Service Unavailable'
assert actual_error.details == {
'message': '{"data": {"key1": "value1", "key2"}',
'status': 'Service Unavailable',
}

View File

@@ -0,0 +1,17 @@
# 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.
#
"""Tests for the Google GenAI SDK's file search stores module."""

View File

@@ -0,0 +1,66 @@
# 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.
#
"""Tests for file_search_stores.create()."""
import pytest
from ... import types
from .. import pytest_helper
# All tests will be run for both Vertex and MLDev.
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name="test_display_name",
parameters=types._CreateFileSearchStoreParameters(
config=types.CreateFileSearchStoreConfig(
display_name="My File Search Store"
)
),
exception_if_vertex="only supported in the Gemini Developer client",
),
pytest_helper.TestTableItem(
name="test_basic",
parameters=types._CreateFileSearchStoreParameters(),
exception_if_vertex="only supported in the Gemini Developer client",
),
]
pytestmark = [
pytest.mark.usefixtures("mock_timestamped_unique_name"),
pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method="file_search_stores.create",
test_table=test_table,
),
]
@pytest.mark.asyncio
async def test_async_display_name(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file_search_store = await client.aio.file_search_stores.create(
config=types.CreateFileSearchStoreConfig(
display_name="My File Search Store"
)
)
@pytest.mark.asyncio
async def test_async_basic(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file_search_store = await client.aio.file_search_stores.create()

View File

@@ -0,0 +1,64 @@
# 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.
#
"""Tests for file_search_stores.delete()."""
import pytest
from ... import types
from .. import pytest_helper
# All tests will be run for both Vertex and MLDev.
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name="test_delete_with_name",
parameters=types._DeleteFileSearchStoreParameters(
name="fileSearchStores/acxjj7m366ln-aw6xyp94icll",
),
exception_if_vertex="only supported in the Gemini Developer client",
),
pytest_helper.TestTableItem(
name="test_delete_with_name_and_force",
parameters=types._DeleteFileSearchStoreParameters(
name="fileSearchStores/7igesc9r2zw9-0mpxpsqubv7s",
config=types.DeleteFileSearchStoreConfig(force=True),
),
exception_if_vertex="only supported in the Gemini Developer client",
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method="file_search_stores.delete",
test_table=test_table,
)
@pytest.mark.asyncio
async def test_async_delete(client):
with pytest_helper.exception_if_vertex(client, ValueError):
await client.aio.file_search_stores.delete(
name="fileSearchStores/my-file-search-store-l65kcyel9lkz"
)
@pytest.mark.asyncio
async def test_async_force_delete(client):
with pytest_helper.exception_if_vertex(client, ValueError):
await client.aio.file_search_stores.delete(
name="fileSearchStores/my-file-search-store-vjtrjw6re8oz",
config=types.DeleteFileSearchStoreConfig(force=True),
)

View File

@@ -0,0 +1,94 @@
# 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.
#
"""Tests for file_search_stores.get()."""
import pytest
from ... import types
from .. import pytest_helper
# A FileSearchStore name known to exist in the test environment.
# Replace with a real one from your test setup.
_EXISTING_FILE_SEARCH_STORE_NAME = "fileSearchStores/5en07ei3kojo-yo8sjqgvx2xf"
_NON_EXISTENT_FILE_SEARCH_STORE_NAME = (
"fileSearchStores/non-existent-file-search-store"
)
_INVALID_FILE_SEARCH_STORE_NAME = (
"fileSearchStores/_invalid_file_search_store_name"
)
_NOT_A_FILE_SEARCH_STORE_NAME = "genai-test-file-search-store"
# All tests will be run for both Vertex and MLDev.
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name="test_get_success",
parameters=types._GetFileSearchStoreParameters(
name=_EXISTING_FILE_SEARCH_STORE_NAME
),
exception_if_vertex="only supported in the Gemini Developer client",
),
pytest_helper.TestTableItem(
name="test_get_not_found",
parameters=types._GetFileSearchStoreParameters(
name=_NON_EXISTENT_FILE_SEARCH_STORE_NAME
),
exception_if_vertex="only supported in the Gemini Developer client",
# Expect an exception to be raised by the mock
exception_if_mldev="PERMISSION_DENIED",
),
pytest_helper.TestTableItem(
name="test_get_invalid_name",
parameters=types._GetFileSearchStoreParameters(
name=_INVALID_FILE_SEARCH_STORE_NAME
),
exception_if_vertex="only supported in the Gemini Developer client",
# Validation should catch this before the API call
exception_if_mldev="INVALID_ARGUMENT",
),
pytest_helper.TestTableItem(
name="test_get_not_a_file_search_store_name",
parameters=types._GetFileSearchStoreParameters(
name=_NOT_A_FILE_SEARCH_STORE_NAME
),
exception_if_vertex="only supported in the Gemini Developer client",
exception_if_mldev="Not Found",
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method="file_search_stores.get",
test_table=test_table,
)
@pytest.mark.asyncio
async def test_async_get(client):
with pytest_helper.exception_if_vertex(client, ValueError):
# This test relies on the mocking/replays set up by pytest_helper
# For a success case:
try:
file_search_store = await client.aio.file_search_stores.get(
name=_EXISTING_FILE_SEARCH_STORE_NAME
)
assert file_search_store is not None
assert file_search_store.name == _EXISTING_FILE_SEARCH_STORE_NAME
except Exception as e:
# Depending on the mock setup, errors might be raised directly
print(f"Async get failed: {e}")
raise

View File

@@ -0,0 +1,112 @@
# 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.
#
"""Tests for file_search_stores.upload_to_file_search_store()."""
import io
import pathlib
import pytest
from ... import types
from .. import pytest_helper
_FILE_SEARCH_STORE_NAME = 'fileSearchStores/my-store-37cbhu1nw16r'
_FILE_NAME = 'files/mk4h34zkv33d'
test_table: list[pytest_helper.TestTableItem] = []
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='filesearchstores.import_file',
test_table=test_table,
)
def test_import(client):
with pytest_helper.exception_if_vertex(client, ValueError):
operation = client.file_search_stores.import_file(
file_search_store_name=_FILE_SEARCH_STORE_NAME, file_name=_FILE_NAME
)
def test_chunking(client):
with pytest_helper.exception_if_vertex(client, ValueError):
operation = client.file_search_stores.import_file(
file_search_store_name=_FILE_SEARCH_STORE_NAME,
file_name=_FILE_NAME,
config=types.ImportFileConfig(
chunking_config=types.ChunkingConfig(
white_space_config=types.WhiteSpaceConfig(
max_tokens_per_chunk=200,
max_overlap_tokens=20,
)
),
),
)
def test_metadata(client):
with pytest_helper.exception_if_vertex(client, ValueError):
operation = client.file_search_stores.import_file(
file_search_store_name=_FILE_SEARCH_STORE_NAME,
file_name=_FILE_NAME,
config=types.ImportFileConfig(
custom_metadata=[
types.CustomMetadata(key='year', numeric_value=2024),
types.CustomMetadata(key='tag', string_value='story'),
],
),
)
@pytest.mark.asyncio
async def test_async_import(client):
with pytest_helper.exception_if_vertex(client, ValueError):
operation = await client.aio.file_search_stores.import_file(
file_search_store_name=_FILE_SEARCH_STORE_NAME,
file_name=_FILE_NAME,
)
@pytest.mark.asyncio
async def test_async_chunking(client):
with pytest_helper.exception_if_vertex(client, ValueError):
operation = await client.aio.file_search_stores.import_file(
file_search_store_name=_FILE_SEARCH_STORE_NAME,
file_name=_FILE_NAME,
config=types.ImportFileConfig(
chunking_config=types.ChunkingConfig(
white_space_config=types.WhiteSpaceConfig(
max_tokens_per_chunk=200,
max_overlap_tokens=20,
)
),
),
)
@pytest.mark.asyncio
async def test_async_metadata(client):
with pytest_helper.exception_if_vertex(client, ValueError):
operation = await client.aio.file_search_stores.import_file(
file_search_store_name=_FILE_SEARCH_STORE_NAME,
file_name=_FILE_NAME,
config=types.ImportFileConfig(
custom_metadata=[
types.CustomMetadata(key='year', numeric_value=2024),
types.CustomMetadata(key='tag', string_value='story'),
],
),
)

View File

@@ -0,0 +1,57 @@
# 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.
#
"""Tests for file_search_stores.list()."""
import pytest
from ... import types
from .. import pytest_helper
# All tests will be run for both Vertex and MLDev.
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name="test_list_default",
parameters=types._ListFileSearchStoresParameters(),
exception_if_vertex="only supported in the Gemini Developer client",
),
pytest_helper.TestTableItem(
name="test_list_with_page_size",
parameters=types._ListFileSearchStoresParameters(
config=types.ListFileSearchStoresConfig(
page_size=5,
),
),
exception_if_vertex="only supported in the Gemini Developer client",
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method="file_search_stores.list",
test_table=test_table,
)
@pytest.mark.asyncio
async def test_async_pager(client):
with pytest_helper.exception_if_vertex(client, ValueError):
# Iterate through the first page.
async for file_search_store in await client.aio.file_search_stores.list(
config={"page_size": 2}
):
assert isinstance(file_search_store, types.FileSearchStore)
break # Only check one item

View File

@@ -0,0 +1,141 @@
# 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.
#
"""Tests for file_search_stores.upload_to_file_search_store()."""
import io
import pathlib
import pytest
from ... import types
from .. import pytest_helper
_FILE_SEARCH_STORE_NAME = 'fileSearchStores/my-store-37cbhu1nw16r'
test_table: list[pytest_helper.TestTableItem] = []
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='filesearchstores.upload_to_file_search_store',
test_table=test_table,
)
def test_file_path(client):
with pytest_helper.exception_if_vertex(client, ValueError):
operation = client.file_search_stores.upload_to_file_search_store(
file_search_store_name=_FILE_SEARCH_STORE_NAME,
file='tests/data/story.pdf',
)
def test_bytesio(client):
with pytest_helper.exception_if_vertex(client, ValueError):
with open('tests/data/story.pdf', 'rb') as f:
with io.BytesIO(f.read()) as buffer:
operation = client.file_search_stores.upload_to_file_search_store(
file_search_store_name=_FILE_SEARCH_STORE_NAME,
file=buffer,
config=types.UploadToFileSearchStoreConfig(
mime_type='application/pdf',
),
)
def test_chunking(client):
with pytest_helper.exception_if_vertex(client, ValueError):
operation = client.file_search_stores.upload_to_file_search_store(
file_search_store_name=_FILE_SEARCH_STORE_NAME,
file='tests/data/story.pdf',
config=types.UploadToFileSearchStoreConfig(
chunking_config=types.ChunkingConfig(
white_space_config=types.WhiteSpaceConfig(
max_tokens_per_chunk=200,
max_overlap_tokens=20,
)
),
),
)
def test_metadata(client):
with pytest_helper.exception_if_vertex(client, ValueError):
operation = client.file_search_stores.upload_to_file_search_store(
file_search_store_name=_FILE_SEARCH_STORE_NAME,
file='tests/data/story.pdf',
config=types.UploadToFileSearchStoreConfig(
custom_metadata=[
types.CustomMetadata(key='year', numeric_value=2024),
types.CustomMetadata(key='tag', string_value='story'),
],
),
)
@pytest.mark.asyncio
async def test_async_file_path(client):
with pytest_helper.exception_if_vertex(client, ValueError):
operation = await client.aio.file_search_stores.upload_to_file_search_store(
file_search_store_name=_FILE_SEARCH_STORE_NAME,
file='tests/data/story.pdf',
)
@pytest.mark.asyncio
async def test_async_bytesio(client):
with pytest_helper.exception_if_vertex(client, ValueError):
with open('tests/data/story.pdf', 'rb') as f:
with io.BytesIO(f.read()) as buffer:
operation = (
await client.aio.file_search_stores.upload_to_file_search_store(
file_search_store_name=_FILE_SEARCH_STORE_NAME,
file=buffer,
config=types.UploadToFileSearchStoreConfig(
mime_type='application/pdf',
),
)
)
@pytest.mark.asyncio
async def test_async_chunking(client):
with pytest_helper.exception_if_vertex(client, ValueError):
operation = await client.aio.file_search_stores.upload_to_file_search_store(
file_search_store_name=_FILE_SEARCH_STORE_NAME,
file='tests/data/story.pdf',
config=types.UploadToFileSearchStoreConfig(
chunking_config=types.ChunkingConfig(
white_space_config=types.WhiteSpaceConfig(
max_tokens_per_chunk=200,
max_overlap_tokens=20,
)
),
),
)
@pytest.mark.asyncio
async def test_async_metadata(client):
with pytest_helper.exception_if_vertex(client, ValueError):
operation = await client.aio.file_search_stores.upload_to_file_search_store(
file_search_store_name=_FILE_SEARCH_STORE_NAME,
file='tests/data/story.pdf',
config=types.UploadToFileSearchStoreConfig(
custom_metadata=[
types.CustomMetadata(key='year', numeric_value=2024),
types.CustomMetadata(key='tag', string_value='story'),
],
),
)

View File

@@ -0,0 +1,17 @@
# 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.
#
"""Tests for the Google GenAI SDK's files module."""

View File

@@ -0,0 +1,46 @@
# 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.
#
"""Test files delete method."""
import pytest
from ... import types
from .. import pytest_helper
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_delete',
parameters=types._DeleteFileParameters(name='files/1g583ke2xdsn'),
exception_if_vertex='only supported in the Gemini Developer client',
skip_in_api_mode=(
'The files have a TTL, they cannot be reliably retrieved for a long'
' time.'
),
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='files.delete',
test_table=test_table,
)
@pytest.mark.asyncio
async def test_async(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = await client.aio.files.get(name='files/n1gls7dyh90q')

View File

@@ -0,0 +1,85 @@
# 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.
#
"""Test files upload method."""
import pathlib
import pytest
from ... import _transformers as t
from ... import types
from .. import pytest_helper
test_table: list[pytest_helper.TestTableItem] = []
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='t.t_file_name',
test_table=test_table,
)
pytest_plugins = ('pytest_asyncio',)
def test_name_transform_name(client):
with pytest_helper.exception_if_vertex(client, ValueError):
for file in client.files.list():
if file.download_uri is not None:
break
else:
raise ValueError('No files found with a `download_uri`.')
file_id = file.name.split('/')[-1]
video = types.Video(uri=file.download_uri)
generated_video = types.GeneratedVideo(video=video)
for f in [
file,
file_id,
file.name,
file.uri,
file.download_uri,
video,
generated_video,
]:
name = t.t_file_name(f)
assert name == file_id
def test_basic_download(client):
with pytest_helper.exception_if_vertex(client, ValueError):
for file in client.files.list():
if file.download_uri is not None:
break
else:
raise ValueError('No files found with a `download_uri`.')
content = client.files.download(file=file)
assert content[4:8] == b'ftyp'
@pytest.mark.asyncio
async def test_basic_download_async(client):
with pytest_helper.exception_if_vertex(client, ValueError):
async for file in await client.aio.files.list():
if file.download_uri is not None:
break
else:
raise ValueError('No files found with a `download_uri`.')
content = await client.aio.files.download(file=file)
assert content[4:8] == b'ftyp'

View File

@@ -0,0 +1,46 @@
# 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.
#
"""Test files get method."""
import pytest
from ... import types
from .. import pytest_helper
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_get',
parameters=types._GetFileParameters(name='files/vjvu9fwk2qj8'),
exception_if_vertex='only supported in the Gemini Developer client',
skip_in_api_mode=(
'The files have a TTL, they cannot be reliably retrieved for a long'
' time.'
),
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='files.get',
test_table=test_table,
)
@pytest.mark.asyncio
async def test_async(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = await client.aio.files.get(name='files/vjvu9fwk2qj8')

View File

@@ -0,0 +1,72 @@
# 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.
#
"""Test files list method."""
import pytest
from ... import types
from .. import pytest_helper
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_not_empty',
exception_if_vertex='only supported in the Gemini Developer client',
parameters=types._ListFilesParameters(
config=types.ListFilesConfig(
page_size=2,
),
),
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='files.list',
test_table=test_table,
)
def test_pager(client):
with pytest_helper.exception_if_vertex(client, ValueError):
files = client.files.list(config={'page_size': 2})
assert 'content-type' in files.sdk_http_response.headers
assert files.name == 'files'
assert files.page_size == 2
assert len(files) <= 2
# Iterate through all the pages. Then next_page() should raise an exception.
for _ in files:
pass
with pytest.raises(IndexError, match='No more pages to fetch.'):
files.next_page()
@pytest.mark.asyncio
async def test_async_pager(client):
with pytest_helper.exception_if_vertex(client, ValueError):
files = await client.aio.files.list(config={'page_size': 2})
assert 'Content-Type' in files.sdk_http_response.headers
assert files.name == 'files'
assert files.page_size == 2
assert len(files) <= 2
# Iterate through all the pages. Then next_page() should raise an exception.
async for _ in files:
pass
with pytest.raises(IndexError, match='No more pages to fetch.'):
await files.next_page()

View File

@@ -0,0 +1,272 @@
# 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.
#
"""Test files register method."""
import json
from unittest import mock
from google.auth import credentials
import httpx
import pytest
from ... import _api_client
from ... import Client
from ... import types
from .. import pytest_helper
class FakeCredentials(credentials.Credentials):
def __init__(self, token="fake_token", expired=False, quota_project_id=None):
super().__init__()
self.token = token
self._expired = expired
self._quota_project_id = quota_project_id
self.refresh_count = 0
@property
def expired(self):
return self._expired
@property
def quota_project_id(self):
return self._quota_project_id
def refresh(self, request):
self.refresh_count += 1
self.token = "refreshed_token"
self._expired = False
@mock.patch.object(_api_client.BaseApiClient, "_request_once", autospec=True)
def test_simple_token(mock_request):
client = Client(api_key="dummy_key")
captured_request = None
def side_effect(self, http_request, stream=False):
nonlocal captured_request
captured_request = http_request
return _api_client.HttpResponse(
headers={},
response_stream=[json.dumps({"files": [{"uri": "files/abc"}]})],
)
mock_request.side_effect = side_effect
with pytest_helper.exception_if_vertex(client, ValueError):
response = client.files.register_files(
auth=FakeCredentials(token="test_token"),
uris=["gs://test-bucket/test-file-1.txt"],
)
assert len(response.files) == 1
assert response.files[0].uri == "files/abc"
assert captured_request.headers["authorization"] == "Bearer test_token"
@mock.patch.object(_api_client.BaseApiClient, "_request_once", autospec=True)
def test_token_refresh(mock_request):
client = Client(api_key="dummy_key")
captured_request = None
def side_effect(self, http_request, stream=False):
nonlocal captured_request
captured_request = http_request
return _api_client.HttpResponse(
headers={},
response_stream=[json.dumps({"files": [{"uri": "files/abc"}]})],
)
mock_request.side_effect = side_effect
with pytest_helper.exception_if_vertex(client, ValueError):
creds = FakeCredentials(expired=True)
response = client.files.register_files(
auth=creds,
uris=["gs://test-bucket/test-file-1.txt"],
)
assert creds.refresh_count == 1
assert len(response.files) == 1
assert response.files[0].uri == "files/abc"
assert captured_request.headers["authorization"] == "Bearer refreshed_token"
@mock.patch.object(_api_client.BaseApiClient, "_request_once", autospec=True)
def test_quota_project(mock_request):
client = Client(api_key="dummy_key")
captured_request = None
def side_effect(self, http_request, stream=False):
nonlocal captured_request
captured_request = http_request
return _api_client.HttpResponse(
headers={},
response_stream=[json.dumps({"files": [{"uri": "files/abc"}]})],
)
mock_request.side_effect = side_effect
with pytest_helper.exception_if_vertex(client, ValueError):
creds = FakeCredentials(quota_project_id="test_project")
response = client.files.register_files(
auth=creds,
uris=["gs://test-bucket/test-file-1.txt"],
)
assert len(response.files) == 1
assert response.files[0].uri == "files/abc"
assert captured_request.headers["x-goog-user-project"] == "test_project"
@mock.patch.object(_api_client.BaseApiClient, "_request_once", autospec=True)
def test_multiple_uris(mock_request):
client = Client(api_key="dummy_key")
def side_effect(self, http_request, stream=False):
return _api_client.HttpResponse(
headers={},
response_stream=[
json.dumps({"files": [{"uri": "files/abc"}, {"uri": "files/def"}]})
],
)
mock_request.side_effect = side_effect
with pytest_helper.exception_if_vertex(client, ValueError):
response = client.files.register_files(
auth=FakeCredentials(),
uris=[
"gs://test-bucket/test-file-1.txt",
"gs://test-bucket/test-file-2.txt",
],
)
assert len(response.files) == 2
assert response.files[0].uri == "files/abc"
assert response.files[1].uri == "files/def"
@pytest.mark.asyncio
@mock.patch.object(
_api_client.BaseApiClient, "_async_request_once", autospec=True
)
async def test_async_single(mock_request):
client = Client(api_key="dummy_key")
async def side_effect(self, http_request, stream=False):
return _api_client.HttpResponse(
headers={},
response_stream=[json.dumps({"files": [{"uri": "files/abc"}]})],
)
mock_request.side_effect = side_effect
with pytest_helper.exception_if_vertex(client, ValueError):
response = await client.aio.files.register_files(
auth=FakeCredentials(),
uris=["gs://test-bucket/test-file-1.txt"],
)
assert len(response.files) == 1
assert response.files[0].uri == "files/abc"
@pytest.mark.asyncio
@mock.patch.object(
_api_client.BaseApiClient, "_async_request_once", autospec=True
)
async def test_async_token_refresh(mock_request):
client = Client(api_key="dummy_key")
captured_request = None
async def side_effect(self, http_request, stream=False):
nonlocal captured_request
captured_request = http_request
return _api_client.HttpResponse(
headers={},
response_stream=[json.dumps({"files": [{"uri": "files/abc"}]})],
)
mock_request.side_effect = side_effect
with pytest_helper.exception_if_vertex(client, ValueError):
creds = FakeCredentials(expired=True)
response = await client.aio.files.register_files(
auth=creds,
uris=["gs://test-bucket/test-file-1.txt"],
)
assert creds.refresh_count == 1
assert len(response.files) == 1
assert response.files[0].uri == "files/abc"
assert captured_request.headers["authorization"] == "Bearer refreshed_token"
@pytest.mark.asyncio
@mock.patch.object(
_api_client.BaseApiClient, "_async_request_once", autospec=True
)
async def test_async_quota_project(mock_request):
client = Client(api_key="dummy_key")
captured_request = None
async def side_effect(self, http_request, stream=False):
nonlocal captured_request
captured_request = http_request
return _api_client.HttpResponse(
headers={},
response_stream=[json.dumps({"files": [{"uri": "files/abc"}]})],
)
mock_request.side_effect = side_effect
with pytest_helper.exception_if_vertex(client, ValueError):
creds = FakeCredentials(quota_project_id="test_project")
response = await client.aio.files.register_files(
auth=creds,
uris=["gs://test-bucket/test-file-1.txt"],
)
assert len(response.files) == 1
assert response.files[0].uri == "files/abc"
assert captured_request.headers["x-goog-user-project"] == "test_project"
@pytest.mark.asyncio
@mock.patch.object(
_api_client.BaseApiClient, "_async_request_once", autospec=True
)
async def test_async_multiple_uris(mock_request):
client = Client(api_key="dummy_key")
async def side_effect(self, http_request, stream=False):
return _api_client.HttpResponse(
headers={},
response_stream=[
json.dumps({"files": [{"uri": "files/abc"}, {"uri": "files/def"}]})
],
)
mock_request.side_effect = side_effect
with pytest_helper.exception_if_vertex(client, ValueError):
response = await client.aio.files.register_files(
auth=FakeCredentials(),
uris=[
"gs://test-bucket/test-file-1.txt",
"gs://test-bucket/test-file-2.txt",
],
)
assert len(response.files) == 2
assert response.files[0].uri == "files/abc"
assert response.files[1].uri == "files/def"

View File

@@ -0,0 +1,70 @@
# 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.
#
"""Test files get method."""
import pytest
from ... import types
from ... import Client
from ... import _api_client
from .. import pytest_helper
import google.auth
# $ gcloud config set project vertex-sdk-dev
# $ gcloud auth application-default login --no-launch-browser --scopes="https://www.googleapis.com/auth/cloud-platform,https://www.googleapis.com/auth/devstorage.read_only"
def get_headers():
try:
credentials, _ = google.auth.default()
token = _api_client.get_token_from_credentials(None, credentials)
headers = {
"Authorization": f"Bearer {token}",}
if credentials.quota_project_id:
headers["x-goog-user-project"] = credentials.quota_project_id
except google.auth.exceptions.DefaultCredentialsError:
# So this can run in replay mode without credentials.
headers = {}
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_register',
parameters=types._InternalRegisterFilesParameters(uris=['gs://unified-genai-dev/image.jpg']),
exception_if_vertex='only supported in the Gemini Developer client',
skip_in_api_mode=(
'The files have a TTL, they cannot be reliably retrieved for a long'
' time.'
),
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='files._register_files',
test_table=test_table,
http_options={
'headers': get_headers(),
},
)
@pytest.mark.asyncio
async def test_async(client):
with pytest_helper.exception_if_vertex(client, ValueError):
files = await client.aio.files._register_files(uris=['gs://unified-genai-dev/image.jpg'])
assert files.files
assert files.files[0].mime_type == 'image/jpeg'

View File

@@ -0,0 +1,255 @@
# 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.
#
"""Test files upload method."""
import io
import pathlib
import pytest
from ... import types
from ... import errors
from .. import pytest_helper
# Upload method is not pydantic.
test_table: list[pytest_helper.TestTableItem] = []
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='files.upload',
test_table=test_table,
)
def test_image_png_upload(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = client.files.upload(file='tests/data/google.png')
assert file.name.startswith('files/')
def test_image_png_upload_with_path(client):
with pytest_helper.exception_if_vertex(client, ValueError):
p = pathlib.Path('tests/data/google.png')
file = client.files.upload(
file=p,
config=types.UploadFileConfig(display_name='test_image_png_path'),
)
assert file.name.startswith('files/')
def test_image_png_upload_with_bytesio(client):
with pytest_helper.exception_if_vertex(client, ValueError):
with open('tests/data/google.png', 'rb') as f:
with io.BytesIO(f.read()) as buffer:
file = client.files.upload(
file=buffer,
config=types.UploadFileConfig(mime_type='image/png'),
)
assert file.name.startswith('files/')
def test_image_png_upload_with_fd(client):
with pytest_helper.exception_if_vertex(client, ValueError):
with open('tests/data/google.png', 'rb') as f:
file = client.files.upload(
file=f,
config=types.UploadFileConfig(mime_type='image/png'),
)
assert file.name.startswith('files/')
def test_image_png_upload_with_config(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = client.files.upload(
file='tests/data/google.png',
config=types.UploadFileConfig(display_name='test_image_png'),
)
assert file.name.startswith('files/')
def test_image_png_upload_with_config_dict(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = client.files.upload(
file='tests/data/google.png', config={'display_name': 'test_image_png'}
)
assert file.name.startswith('files/')
def test_image_jpg_upload(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = client.files.upload(file='tests/data/google.jpg')
assert file.name.startswith('files/')
def test_image_jpg_upload_with_config(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = client.files.upload(
file='tests/data/google.jpg',
config=types.UploadFileConfig(display_name='test_image_jpg'),
)
assert file.name.startswith('files/')
def test_image_jpg_upload_with_config_dict(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = client.files.upload(
file='tests/data/google.jpg', config={'display_name': 'test_image_jpg'}
)
assert file.name.startswith('files/')
def test_application_pdf_file_upload(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = client.files.upload(file='tests/data/story.pdf')
assert file.name.startswith('files/')
def test_application_pdf_upload_with_config(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = client.files.upload(
file='tests/data/story.pdf',
config=types.UploadFileConfig(display_name='test_application_pdf'),
)
assert file.name.startswith('files/')
def test_application_pdf_upload_with_config_dict(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = client.files.upload(
file='tests/data/story.pdf',
config={'display_name': 'test_application_pdf'},
)
assert file.name.startswith('files/')
def test_video_mp4_file_upload(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = client.files.upload(file='tests/data/animal.mp4')
assert file.name.startswith('files/')
def test_video_mp4_upload_with_config(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = client.files.upload(
file='tests/data/animal.mp4',
config=types.UploadFileConfig(display_name='test_video_mp4'),
)
assert file.name.startswith('files/')
def test_video_mp4_upload_with_config_dict(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = client.files.upload(
file='tests/data/animal.mp4', config={'display_name': 'test_video_mp4'}
)
assert file.name.startswith('files/')
def test_audio_m4a_file_upload(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = client.files.upload(
file='tests/data/pixel.m4a',
config=types.UploadFileConfig(mime_type='audio/mp4'),
)
assert file.name.startswith('files/')
def test_audio_m4a_upload_with_config(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = client.files.upload(
file='tests/data/pixel.m4a',
config=types.UploadFileConfig(
display_name='test_audio_m4a', mime_type='audio/mp4'
),
)
assert file.name.startswith('files/')
def test_audio_m4a_upload_with_config_dict(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = client.files.upload(
file='tests/data/pixel.m4a',
config={'display_name': 'test_audio_m4a', 'mime_type': 'audio/mp4'},
)
assert file.name.startswith('files/')
def test_bad_mime_type(client):
with pytest_helper.exception_if_vertex(client, ValueError):
with pytest.raises(errors.APIError, match="Unsupported MIME"):
file = client.files.upload(
file=io.BytesIO(b'test'),
config={'mime_type': 'bad/mime_type'},
)
@pytest.mark.asyncio
async def test_image_upload_async(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = await client.aio.files.upload(file='tests/data/google.png')
assert file.name.startswith('files/')
@pytest.mark.asyncio
async def test_image_upload_with_config_async(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = await client.aio.files.upload(
file='tests/data/google.png',
config=types.UploadFileConfig(display_name='test_image'),
)
assert file.name.startswith('files/')
@pytest.mark.asyncio
async def test_image_upload_with_config_dict_async(client):
with pytest_helper.exception_if_vertex(client, ValueError):
file = await client.aio.files.upload(
file='tests/data/google.png',
config={
'display_name': 'test_image',
'http_options': {'timeout': '8000'},
},
)
assert file.name.startswith('files/')
@pytest.mark.asyncio
async def test_image_upload_with_bytesio_async(client):
with pytest_helper.exception_if_vertex(client, ValueError):
with open('tests/data/google.png', 'rb') as f:
buffer = io.BytesIO(f.read())
file = await client.aio.files.upload(
file=buffer,
config=types.UploadFileConfig(
mime_type='image/png'),
)
assert file.name.startswith('files/')
@pytest.mark.asyncio
async def test_unknown_path_upload_async(client):
with pytest_helper.exception_if_vertex(client, ValueError):
try:
await client.aio.files.upload(file='unknown_path')
except FileNotFoundError as e:
assert 'is not a valid file path' in str(e)
@pytest.mark.asyncio
async def test_bad_mime_type_async(client):
with pytest_helper.exception_if_vertex(client, ValueError):
with pytest.raises(errors.APIError, match="Unsupported MIME"):
file = await client.aio.files.upload(
file=io.BytesIO(b'test'),
config={'mime_type': 'bad/mime_type'},
)

View File

@@ -0,0 +1,28 @@
# 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 os
import pytest
IS_NOT_GITHUB_ACTIONS = os.getenv('GITHUB_ACTIONS') != 'true'
@pytest.mark.skipif(IS_NOT_GITHUB_ACTIONS,
reason='This test is only run on GitHub Actions.')
def test_library_can_be_imported_without_optional_dependencies():
"""Tests that the library can be imported without optional dependencies.
"""
from google import genai
from google.genai import types

View File

@@ -0,0 +1,476 @@
# 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.
#
"""Tests for Interactions API."""
from ... import Client
from ... import _base_url
from unittest import mock
import pytest
from httpx import Request, Response
from ..._api_client import AsyncHttpxClient, BaseApiClient
from httpx import Client as HTTPClient
import os
ENV_VARS = [
"GOOGLE_CLOUD_PROJECT",
"GEMINI_API_KEY",
"GOOGLE_API_KEY",
"GOOGLE_CLOUD_LOCATION",
]
@pytest.fixture(autouse=True)
def clear_env_vars(monkeypatch):
for var in ENV_VARS:
monkeypatch.delenv(var, raising=False)
def test_interactions_gemini_url(monkeypatch):
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
client = Client()
with mock.patch.object(HTTPClient, "send") as mock_send:
mock_send.return_value = Response(200, request=Request('POST', ''))
client.interactions.create(
model='gemini-1.5-flash',
input='Hello',
)
mock_send.assert_called_once()
request = mock_send.call_args[0][0]
assert str(request.url).endswith('/v1beta/interactions')
assert request.headers['x-goog-api-key'] == 'test-api-key'
def test_interactions_gemini_no_vertex_auth(monkeypatch):
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
client = Client()
with (
mock.patch.object(BaseApiClient, "_access_token") as mock_access_token,
mock.patch.object(HTTPClient, "send") as mock_send,
):
mock_send.return_value = Response(200, request=Request('POST', ''))
client.interactions.create(
model='gemini-1.5-flash',
input='Hello',
)
mock_access_token.assert_not_called()
def test_interactions_gemini_retry(monkeypatch):
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
client = Client()
client._api_client.max_retries = 2
with mock.patch.object(HTTPClient, "send") as mock_send:
mock_send.side_effect = [
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1"}),
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1"}),
Response(200, request=Request('POST', '')),
]
client.interactions.create(model='gemini-1.5-flash', input='Hello')
assert mock_send.call_count == 3
def test_interactions_gemini_extra_headers(monkeypatch):
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
client = Client()
with mock.patch.object(HTTPClient, "send") as mock_send:
mock_send.return_value = Response(200, request=Request('POST', ''))
client.interactions.create(
model='gemini-1.5-flash',
input='Hello',
extra_headers={'X-Custom-Header': 'TestValue'}
)
mock_send.assert_called_once()
request = mock_send.call_args[0][0]
assert request.headers['x-custom-header'] == 'TestValue'
assert request.headers['x-goog-api-key'] == 'test-api-key'
def test_interactions_vertex_auth_header():
from ..._api_client import BaseApiClient
from ..._interactions._base_client import SyncAPIClient
from httpx import Client as HTTPClient
creds = mock.Mock()
creds.quota_project_id = "test-quota-project"
client = Client(vertexai=True, project='test-project', location='us-central1', credentials=creds)
with (
mock.patch.object(
BaseApiClient, "_access_token", return_value='fake-vertex-token'
) as mock_access_token,
mock.patch.object(
HTTPClient, "send",
return_value=mock.Mock(),
) as mock_send,
):
response = client.interactions.create(
model='gemini-2.5-flash',
input='What is the largest planet in our solar system?',
)
mock_send.assert_called_once()
mock_access_token.assert_called_once()
args, kwargs = mock_send.call_args
headers = args[0].headers
assert any(
key == "authorization" and value == 'Bearer fake-vertex-token'
for key, value in headers.items())
assert any(
key == "x-goog-user-project" and value == 'test-quota-project'
for key, value in headers.items())
def test_interactions_vertex_key_no_auth_header():
from ..._api_client import BaseApiClient
from httpx import Client as HTTPClient
creds = mock.Mock()
client = Client(vertexai=True, api_key='test-api-key')
with (
mock.patch.object(
BaseApiClient, "_access_token", return_value='fake-vertex-token'
) as mock_access_token,
mock.patch.object(
HTTPClient, "send",
return_value=mock.Mock(),
) as mock_send,
):
response = client.interactions.create(
model='gemini-2.5-flash',
input='What is the largest planet in our solar system?',
)
mock_send.assert_called_once()
mock_access_token.assert_not_called()
args, kwargs = mock_send.call_args
headers = args[0].headers
assert any(
key == "x-goog-api-key" and value == 'test-api-key'
for key, value in headers.items())
def test_interactions_vertex_url():
creds = mock.Mock()
creds.quota_project_id = "test-quota-project"
client = Client(vertexai=True, project='test-project', location='us-central1', credentials=creds)
with mock.patch("httpx.Client.send") as mock_send:
mock_send.return_value = Response(200, request=Request('POST', ''))
client.interactions.create(
model='gemini-1.5-flash',
input='Hello',
)
mock_send.assert_called_once()
request = mock_send.call_args[0][0]
assert str(request.url) == 'https://us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/interactions'
def test_interactions_vertex_auth_refresh_on_retry():
from ..._api_client import BaseApiClient
from httpx import Client as HTTPClient
creds = mock.Mock()
creds.quota_project_id = "test-quota-project"
client = Client(vertexai=True, project='test-project', location='us-central1', credentials=creds)
client._api_client.max_retries = 2
token_values = ['token1', 'token2', 'token3']
token_iter = iter(token_values)
def get_token():
return next(token_iter)
with (
mock.patch.object(BaseApiClient, "_access_token", side_effect=get_token) as mock_access_token,
mock.patch.object(HTTPClient, "send") as mock_send,
):
mock_send.side_effect = [
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1"}),
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1"}),
Response(200, request=Request('POST', '')),
]
client.interactions.create(model='gemini-1.5-flash', input='Hello')
assert mock_access_token.call_count == 3
assert mock_send.call_count == 3
# Check headers of each call
for i in range(3):
headers = mock_send.call_args_list[i][0][0].headers
assert headers['authorization'] == f'Bearer {token_values[i]}'
def test_interactions_vertex_extra_headers_override():
from ..._api_client import BaseApiClient
from httpx import Client as HTTPClient
creds = mock.Mock()
creds.quota_project_id = "test-quota-project"
client = Client(vertexai=True, project='test-project', location='us-central1', credentials=creds)
with (
mock.patch.object(BaseApiClient, "_access_token", return_value='default-token') as mock_access_token,
mock.patch.object(HTTPClient, "send") as mock_send,
):
mock_send.return_value = Response(200, request=Request('POST', ''))
# Override Authorization
client.interactions.create(
model='gemini-1.5-flash',
input='Hello',
extra_headers={'Authorization': 'Bearer manual-token'}
)
mock_send.assert_called_once()
headers = mock_send.call_args[0][0].headers
assert headers['authorization'] == 'Bearer manual-token'
mock_access_token.assert_not_called() # Should not fetch default token
mock_send.reset_mock()
mock_access_token.reset_mock()
# Provide API Key
client.interactions.create(
model='gemini-1.5-flash',
input='Hello',
extra_headers={'x-goog-api-key': 'manual-key'}
)
mock_send.assert_called_once()
headers = mock_send.call_args[0][0].headers
assert headers['x-goog-api-key'] == 'manual-key'
assert 'authorization' not in headers
mock_access_token.assert_not_called()
@pytest.mark.asyncio
async def test_async_interactions_gemini_url(monkeypatch):
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
client = Client()
with mock.patch.object(AsyncHttpxClient, "send") as mock_send:
mock_send.return_value = Response(200, request=Request('POST', ''))
await client.aio.interactions.create(
model='gemini-1.5-flash',
input='Hello',
)
mock_send.assert_called_once()
request = mock_send.call_args[0][0]
assert str(request.url).endswith('/v1beta/interactions')
assert request.headers['x-goog-api-key'] == 'test-api-key'
@pytest.mark.asyncio
async def test_async_interactions_gemini_no_vertex_auth(monkeypatch):
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
client = Client()
with (
mock.patch.object(BaseApiClient, "_async_access_token") as mock_access_token,
mock.patch.object(AsyncHttpxClient, "send") as mock_send,
):
mock_send.return_value = Response(200, request=Request('POST', ''))
await client.aio.interactions.create(
model='gemini-1.5-flash',
input='Hello',
)
mock_access_token.assert_not_called()
@pytest.mark.asyncio
async def test_async_interactions_gemini_retry(monkeypatch):
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
client = Client()
client.aio._api_client.max_retries = 2
with mock.patch.object(AsyncHttpxClient, "send") as mock_send:
mock_send.side_effect = [
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1"}),
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1"}),
Response(200, request=Request('POST', '')),
]
await client.aio.interactions.create(model='gemini-1.5-flash', input='Hello')
assert mock_send.call_count == 3
@pytest.mark.asyncio
async def test_async_interactions_gemini_extra_headers(monkeypatch):
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
client = Client()
with mock.patch.object(AsyncHttpxClient, "send") as mock_send:
mock_send.return_value = Response(200, request=Request('POST', ''))
await client.aio.interactions.create(
model='gemini-1.5-flash',
input='Hello',
extra_headers={'X-Custom-Header': 'TestValue'}
)
mock_send.assert_called_once()
request = mock_send.call_args[0][0]
assert request.headers['x-custom-header'] == 'TestValue'
assert request.headers['x-goog-api-key'] == 'test-api-key'
@pytest.mark.asyncio
async def test_async_interactions_vertex_auth_header():
from ..._api_client import BaseApiClient
from ..._interactions._base_client import SyncAPIClient
from ..._api_client import AsyncHttpxClient
creds = mock.Mock()
creds.quota_project_id = "test-quota-project"
client = Client(vertexai=True, project='test-project', location='us-central1', credentials=creds)
with (
mock.patch.object(
BaseApiClient, "_async_access_token", return_value='fake-vertex-token'
) as mock_access_token,
mock.patch.object(
AsyncHttpxClient, "send",
return_value=mock.Mock(),
) as mock_send,
):
response = await client.aio.interactions.create(
model='gemini-2.5-flash',
input='What is the largest planet in our solar system?',
)
mock_send.assert_called_once()
mock_access_token.assert_called_once()
args, kwargs = mock_send.call_args
headers = args[0].headers
assert any(
key == "authorization" and value == 'Bearer fake-vertex-token'
for key, value in headers.items())
assert any(
key == "x-goog-user-project" and value == 'test-quota-project'
for key, value in headers.items())
@pytest.mark.asyncio
async def test_async_interactions_vertex_key_no_auth_header():
from ..._api_client import BaseApiClient
client = Client(vertexai=True, api_key='test-api-key')
with (
mock.patch.object(
BaseApiClient, "_async_access_token", return_value='fake-vertex-token'
) as mock_access_token,
mock.patch.object(
AsyncHttpxClient, "send",
return_value=mock.Mock(),
) as mock_send,
):
response = await client.aio.interactions.create(
model='gemini-2.5-flash',
input='What is the largest planet in our solar system?',
)
mock_send.assert_called_once()
mock_access_token.assert_not_called()
args, kwargs = mock_send.call_args
headers = args[0].headers
assert any(
key == "x-goog-api-key" and value == 'test-api-key'
for key, value in headers.items())
@pytest.mark.asyncio
async def test_async_interactions_vertex_url():
from ..._api_client import AsyncHttpxClient
creds = mock.Mock()
creds.quota_project_id = "test-quota-project"
client = Client(vertexai=True, project='test-project', location='us-central1', credentials=creds)
with mock.patch.object(AsyncHttpxClient, "send") as mock_send:
mock_send.return_value = Response(200, request=Request('POST', ''))
await client.aio.interactions.create(
model='gemini-1.5-flash',
input='Hello',
)
mock_send.assert_called_once()
request = mock_send.call_args[0][0]
assert str(request.url) == 'https://us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/interactions'
@pytest.mark.asyncio
async def test_async_interactions_vertex_auth_refresh_on_retry():
from ..._api_client import BaseApiClient
from ..._api_client import AsyncHttpxClient
creds = mock.Mock()
creds.quota_project_id = "test-quota-project"
client = Client(vertexai=True, project='test-project', location='us-central1', credentials=creds)
client.aio._api_client.max_retries = 2
token_values = ['token1', 'token2', 'token3']
token_iter = iter(token_values)
async def get_token():
return next(token_iter)
with (
mock.patch.object(BaseApiClient, "_async_access_token", side_effect=get_token) as mock_access_token,
mock.patch.object(AsyncHttpxClient, "send") as mock_send,
):
mock_send.side_effect = [
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1"}),
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1"}),
Response(200, request=Request('POST', '')),
]
await client.aio.interactions.create(model='gemini-1.5-flash', input='Hello')
assert mock_access_token.call_count == 3
assert mock_send.call_count == 3
for i in range(3):
headers = mock_send.call_args_list[i][0][0].headers
assert headers['authorization'] == f'Bearer {token_values[i]}'
@pytest.mark.asyncio
async def test_async_interactions_vertex_extra_headers_override():
from ..._api_client import BaseApiClient
from ..._api_client import AsyncHttpxClient
creds = mock.Mock()
creds.quota_project_id = "test-quota-project"
client = Client(vertexai=True, project='test-project', location='us-central1', credentials=creds)
with (
mock.patch.object(BaseApiClient, "_async_access_token", return_value='default-token') as mock_access_token,
mock.patch.object(AsyncHttpxClient, "send") as mock_send,
):
mock_send.return_value = Response(200, request=Request('POST', ''))
# Override Authorization
await client.aio.interactions.create(
model='gemini-1.5-flash',
input='Hello',
extra_headers={'Authorization': 'Bearer manual-token'}
)
mock_send.assert_called_once()
headers = mock_send.call_args[0][0].headers
assert headers['authorization'] == 'Bearer manual-token'
mock_access_token.assert_not_called()
mock_send.reset_mock()
mock_access_token.reset_mock()
# Provide API Key
await client.aio.interactions.create(
model='gemini-1.5-flash',
input='Hello',
extra_headers={'x-goog-api-key': 'manual-key'}
)
mock_send.assert_called_once()
headers = mock_send.call_args[0][0].headers
assert headers['x-goog-api-key'] == 'manual-key'
assert 'authorization' not in headers
mock_access_token.assert_not_called()

View File

@@ -0,0 +1,84 @@
# 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.
from unittest import mock
import pytest
from ... import client as client_lib
pytest_plugins = ("pytest_asyncio",)
def test_client_future_warning():
with mock.patch.object(
client_lib, "_interactions_experimental_warned", new=False
):
client = client_lib.Client(
api_key="placeholder",
http_options={
"api_version": "v1alpha",
}
)
with pytest.warns(
UserWarning, match="Interactions.*experimental"
):
_ = client.interactions
def test_client_timeout():
with mock.patch.object(
client_lib, "GeminiNextGenAPIClient", spec_set=True
) as mock_nextgen_client:
client = client_lib.Client(
api_key="placeholder",
http_options={"api_version": "v1alpha", "timeout": 5000},
)
_ = client.interactions
mock_nextgen_client.assert_called_once_with(
base_url=mock.ANY,
api_key="placeholder",
api_version="v1alpha",
default_headers=mock.ANY,
http_client=mock.ANY,
timeout=5.0,
max_retries=mock.ANY,
client_adapter=mock.ANY,
)
@pytest.mark.asyncio
async def test_async_client_timeout():
with mock.patch.object(
client_lib, "AsyncGeminiNextGenAPIClient", spec_set=True
) as mock_nextgen_client:
client = client_lib.Client(
api_key="placeholder",
http_options={"api_version": "v1alpha", "timeout": 5000},
)
_ = client.aio.interactions
mock_nextgen_client.assert_called_once_with(
base_url=mock.ANY,
api_key="placeholder",
api_version="v1alpha",
default_headers=mock.ANY,
http_client=mock.ANY,
timeout=5.0,
max_retries=mock.ANY,
client_adapter=mock.ANY,
)

View File

@@ -0,0 +1,105 @@
# 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.
#
"""Tests for Interactions API URL paths."""
from unittest import mock
import pytest
from httpx import Request, Response
from ..._api_client import AsyncHttpxClient
from httpx import Client as HTTPClient
from .. import pytest_helper
import google.auth
@mock.patch.object(google.auth, "default", autospec=True)
def test_interactions_paths(mock_auth_default, client):
interaction_id = "test-interaction-id"
mock_creds = mock.Mock()
mock_creds.token = "test-token"
mock_creds.expired = False
mock_creds.quota_project_id = "test-quota-project"
mock_auth_default.return_value = (mock_creds, "test-project")
if client._api_client.vertexai:
expected_base_url = f'https://{client._api_client.location}-aiplatform.googleapis.com/v1beta1/projects/{client._api_client.project}/locations/{client._api_client.location}'
else:
expected_base_url = "https://generativelanguage.googleapis.com/v1beta"
with mock.patch.object(HTTPClient, "send") as mock_send:
mock_send.return_value = Response(200, request=Request('GET', ''))
client.interactions.get(id=interaction_id)
mock_send.assert_called_once()
request = mock_send.call_args[0][0]
assert str(request.url) == f'{expected_base_url}/interactions/{interaction_id}'
mock_send.reset_mock()
mock_send.return_value = Response(200, request=Request('POST', ''))
client.interactions.cancel(id=interaction_id)
mock_send.assert_called_once()
request = mock_send.call_args[0][0]
assert str(request.url) == f'{expected_base_url}/interactions/{interaction_id}/cancel'
mock_send.reset_mock()
mock_send.return_value = Response(200, request=Request('DELETE', ''))
client.interactions.delete(id=interaction_id)
mock_send.assert_called_once()
request = mock_send.call_args[0][0]
assert str(request.url) == f'{expected_base_url}/interactions/{interaction_id}'
@pytest.mark.asyncio
@mock.patch.object(google.auth, "default", autospec=True)
async def test_async_interactions_paths(mock_auth_default, client):
interaction_id = "test-interaction-id"
mock_creds = mock.Mock()
mock_creds.token = "test-token"
mock_creds.expired = False
mock_creds.quota_project_id = "test-quota-project"
mock_auth_default.return_value = (mock_creds, "test-project")
if client._api_client.vertexai:
expected_base_url = f'https://{client._api_client.location}-aiplatform.googleapis.com/v1beta1/projects/{client._api_client.project}/locations/{client._api_client.location}'
else:
expected_base_url = "https://generativelanguage.googleapis.com/v1beta"
with mock.patch.object(AsyncHttpxClient, "send") as mock_send:
mock_send.return_value = Response(200, request=Request('GET', ''))
await client.aio.interactions.get(id=interaction_id)
mock_send.assert_called_once()
request = mock_send.call_args[0][0]
assert str(request.url) == f'{expected_base_url}/interactions/{interaction_id}'
mock_send.reset_mock()
mock_send.return_value = Response(200, request=Request('POST', ''))
await client.aio.interactions.cancel(id=interaction_id)
mock_send.assert_called_once()
request = mock_send.call_args[0][0]
assert str(request.url) == f'{expected_base_url}/interactions/{interaction_id}/cancel'
mock_send.reset_mock()
mock_send.return_value = Response(200, request=Request('DELETE', ''))
await client.aio.interactions.delete(id=interaction_id)
mock_send.assert_called_once()
request = mock_send.call_args[0][0]
assert str(request.url) == f'{expected_base_url}/interactions/{interaction_id}'
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_table=[],
)

View File

@@ -0,0 +1,16 @@
# 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.
#

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,362 @@
# 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.
#
"""Tests for live_music.py."""
import contextlib
import json
from typing import AsyncIterator
from unittest import mock
from unittest.mock import AsyncMock
from unittest.mock import Mock
from unittest.mock import patch
import warnings
from google.oauth2.credentials import Credentials
import pytest
from websockets import client
from ... import _api_client as api_client
from ... import _common
from ... import Client
from ... import client as gl_client
from ... import live
from ... import live_music
from ... import types
from .. import pytest_helper
try:
import aiohttp
AIOHTTP_NOT_INSTALLED = False
except ImportError:
AIOHTTP_NOT_INSTALLED = True
aiohttp = mock.MagicMock()
requires_aiohttp = pytest.mark.skipif(
AIOHTTP_NOT_INSTALLED, reason="aiohttp is not installed, skipping test."
)
def mock_api_client(vertexai=False, credentials=None):
api_client = mock.MagicMock(spec=gl_client.BaseApiClient)
if not vertexai:
api_client.api_key = 'TEST_API_KEY'
api_client.location = None
api_client.project = None
else:
api_client.api_key = None
api_client.location = 'us-central1'
api_client.project = 'test_project'
api_client._host = lambda: 'test_host'
api_client._credentials = credentials
api_client._http_options = types.HttpOptions.model_validate(
{'headers': {}}
) # Ensure headers exist
api_client.vertexai = vertexai
api_client._api_client = api_client
return api_client
@pytest.fixture
def mock_websocket():
websocket = AsyncMock(spec=client.ClientConnection)
websocket.send = AsyncMock()
websocket.recv = AsyncMock(
return_value=b"""{
"serverContent": {
"audioChunks": [
{
"data": "Z2VsYmFuYW5h",
"mimeType": "audio/l16;rate=48000;channels=2",
"sourceMetadata": {
"clientContent": {
"weightedPrompts": [
{
"text": "Jazz",
"weight": 1
}
]
},
"musicGenerationConfig": {
"seed": -957124937,
"bpm": 140,
"scale": "A_FLAT_MAJOR_F_MINOR"
}
}
}
]
}
}"""
) # Default response
websocket.close = AsyncMock()
return websocket
async def get_connect_message(api_client, model):
mock_ws = AsyncMock()
mock_ws.send = AsyncMock()
mock_ws.recv = AsyncMock(return_value=b'some response')
mock_google_auth_default = Mock(return_value=(None, None))
mock_creds = Mock(token='test_token')
mock_google_auth_default.return_value = (mock_creds, None)
@contextlib.asynccontextmanager
async def mock_connect(uri, additional_headers=None):
yield mock_ws
@patch('google.auth.default', new=mock_google_auth_default)
@patch.object(live_music, 'connect', new=mock_connect)
async def _test_connect():
live_module = live.AsyncLive(api_client)
async with live_module.music.connect(
model=model,
):
pass
mock_ws.send.assert_called_once()
return json.loads(mock_ws.send.call_args[0][0])
return await _test_connect()
def test_mldev_from_env(monkeypatch):
api_key = 'google_api_key'
monkeypatch.setenv('GOOGLE_API_KEY', api_key)
client = Client()
assert not client.aio.live.music._api_client.vertexai
assert client.aio.live.music._api_client.api_key == api_key
assert isinstance(client.aio.live._api_client, api_client.BaseApiClient)
@requires_aiohttp
def test_vertex_from_env(monkeypatch):
project_id = 'fake_project_id'
location = 'fake-location'
monkeypatch.setenv('GOOGLE_GENAI_USE_VERTEXAI', 'true')
monkeypatch.setenv('GOOGLE_CLOUD_PROJECT', project_id)
monkeypatch.setenv('GOOGLE_CLOUD_LOCATION', location)
client = Client()
assert client.aio.live.music._api_client.vertexai
assert client.aio.live.music._api_client.project == project_id
assert isinstance(client.aio.live._api_client, api_client.BaseApiClient)
def test_websocket_base_url():
base_url = 'https://test.com'
api_client = gl_client.BaseApiClient(
api_key='google_api_key',
http_options={'base_url': base_url},
)
assert api_client._websocket_base_url() == 'wss://test.com'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_send_weighted_prompts(
mock_websocket, vertexai
):
session = live_music.AsyncMusicSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
if vertexai:
with pytest.raises(NotImplementedError):
await session.set_weighted_prompts(prompts=[types.WeightedPrompt(text='Jazz', weight=1)])
return
await session.set_weighted_prompts(prompts=[types.WeightedPrompt(text='Jazz', weight=1)])
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'clientContent' in sent_data
assert sent_data['clientContent']['weightedPrompts'][0]['text'] == 'Jazz'
assert sent_data['clientContent']['weightedPrompts'][0]['weight'] == 1
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_send_config(
mock_websocket, vertexai
):
session = live_music.AsyncMusicSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
if vertexai:
with pytest.raises(NotImplementedError):
await session.set_music_generation_config(
config=types.LiveMusicGenerationConfig(
bpm=140,
music_generation_mode=types.MusicGenerationMode.VOCALIZATION,
)
)
return
await session.set_music_generation_config(
config=types.LiveMusicGenerationConfig(
bpm=140,
music_generation_mode=types.MusicGenerationMode.VOCALIZATION,
)
)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'musicGenerationConfig' in sent_data
assert sent_data['musicGenerationConfig']['bpm'] == 140
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_control_signal_play(
mock_websocket, vertexai
):
session = live_music.AsyncMusicSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
if vertexai:
with pytest.raises(NotImplementedError):
await session.play()
return
await session.play()
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'playbackControl' in sent_data
assert 'PLAY' in sent_data['playbackControl']
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_control_signal_pause(
mock_websocket, vertexai
):
session = live_music.AsyncMusicSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
if vertexai:
with pytest.raises(NotImplementedError):
await session.pause()
return
await session.pause()
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'playbackControl' in sent_data
assert 'PAUSE' in sent_data['playbackControl']
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_control_signal_stop(
mock_websocket, vertexai
):
session = live_music.AsyncMusicSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
if vertexai:
with pytest.raises(NotImplementedError):
await session.stop()
return
await session.stop()
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'playbackControl' in sent_data
assert 'STOP' in sent_data['playbackControl']
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_control_signal_reset_context(
mock_websocket, vertexai
):
session = live_music.AsyncMusicSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
if vertexai:
with pytest.raises(NotImplementedError):
await session.reset_context()
return
await session.reset_context()
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'playbackControl' in sent_data
assert 'RESET_CONTEXT' in sent_data['playbackControl']
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_receive( mock_websocket, vertexai):
session = live_music.AsyncMusicSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
if vertexai:
with pytest.raises(NotImplementedError):
async for _ in session.receive():
pass
return
async for response in session.receive():
assert isinstance(response, types.LiveMusicServerMessage)
audio_chunk = response.server_content.audio_chunks[0]
# Data contains decoded b64 audio
assert audio_chunk.data == b'gelbanana'
assert audio_chunk.mime_type == 'audio/l16;rate=48000;channels=2'
assert audio_chunk.source_metadata.client_content.weighted_prompts[0].text == 'Jazz'
assert audio_chunk.source_metadata.client_content.weighted_prompts[0].weight == 1
assert audio_chunk.source_metadata.music_generation_config.bpm == 140
break
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_receive_error(
mock_websocket, vertexai
):
mock_websocket.recv = AsyncMock(return_value='invalid json')
session = live_music.AsyncMusicSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
with pytest.raises(ValueError):
await session.receive().__anext__()
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_close( mock_websocket, vertexai):
session = live_music.AsyncMusicSession(
mock_api_client(vertexai=vertexai), mock_websocket
)
await session.close()
mock_websocket.close.assert_called_once()
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_setup_to_api(vertexai):
if vertexai:
with pytest.raises(NotImplementedError):
await get_connect_message(
mock_api_client(vertexai=vertexai),
model='test_model'
)
return
result = await get_connect_message(
mock_api_client(vertexai=vertexai),
model='test_model'
)
expected_result = {'setup': {}}
if vertexai:
# Vertex is not supported yet
assert False
else:
expected_result['setup']['model'] = 'models/test_model'
assert result == expected_result

View File

@@ -0,0 +1,163 @@
# 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.
#
"""Tests for live response handling."""
import json
from typing import cast
from unittest import mock
from unittest.mock import AsyncMock
import pytest
from ... import _api_client as api_client
from ... import _common
from ... import Client
from ... import client as gl_client
from ... import live
from ... import types
def mock_api_client(vertexai=False):
"""Creates a mock BaseApiClient."""
mock_client = mock.MagicMock(spec=gl_client.BaseApiClient)
if not vertexai:
mock_client.api_key = 'TEST_API_KEY'
mock_client.location = None
mock_client.project = None
else:
mock_client.api_key = None
mock_client.location = 'us-central1'
mock_client.project = 'test_project'
mock_client._host = lambda: 'test_host'
mock_client._http_options = types.HttpOptions.model_validate(
{'headers': {}}
)
mock_client.vertexai = vertexai
return mock_client
@pytest.fixture
def mock_websocket():
"""Provides a mock websocket connection."""
# Use live.ClientConnection if that's the specific type hint in AsyncSession
websocket = AsyncMock(spec=live.ClientConnection)
websocket.send = AsyncMock()
# Set default recv value, will be overridden in the test
websocket.recv = AsyncMock(return_value='{}')
websocket.close = AsyncMock()
return websocket
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_receive_server_content(mock_websocket, vertexai):
raw_response_json = json.dumps({
"usageMetadata": {
"promptTokenCount": 15,
"responseTokenCount": 25,
"candidatesTokenCount": 50,
"totalTokenCount": 200,
"responseTokensDetails": [
{
"tokenCount": 20,
"modality": "TEXT",
}
],
"candidatesTokensDetails": [
{
"tokenCount": 10,
"modality": "TEXT",
}
],
},
"serverContent": {
"modelTurn": {
"parts": [{"text": "This is a simple response."}]
},
"turnComplete": True,
"groundingMetadata": {
"web_search_queries": ["test query"],
"groundingChunks": [{
"web": {
"domain": "google.com",
"title": "Search results",
}
}]
}
}
})
mock_websocket.recv.return_value = raw_response_json
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
result = await session._receive()
# Assert the results
assert isinstance(result, types.LiveServerMessage)
assert (
result.server_content.model_turn.parts[0].text
== "This is a simple response."
)
assert result.server_content.turn_complete
assert result.server_content.grounding_metadata.web_search_queries == ["test query"]
assert result.server_content.grounding_metadata.grounding_chunks[0].web.domain == "google.com"
assert result.server_content.grounding_metadata.grounding_chunks[0].web.title == "Search results"
# Verify usageMetadata was parsed
assert isinstance(result.usage_metadata, types.UsageMetadata)
assert result.usage_metadata.prompt_token_count == 15
assert result.usage_metadata.total_token_count == 200
if not vertexai:
assert result.usage_metadata.response_token_count == 25
assert result.usage_metadata.response_tokens_details[0].token_count == 20
else:
# VertexAI maps candidatesTokenCount to responseTokenCount and maps
# candidatesTokensDetails to responseTokensDetails.
assert result.usage_metadata.response_token_count == 50
assert result.usage_metadata.response_tokens_details[0].token_count == 10
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_receive_server_content_with_turn_reason(mock_websocket, vertexai):
"""Tests parsing of LiveServerContent with turn_complete_reason and waiting_for_input."""
raw_response_json = json.dumps({
"serverContent": {
"modelTurn": {
"parts": [{"text": "Please provide more details."}]
},
"turnComplete": True,
"turnCompleteReason": "NEED_MORE_INPUT",
"waitingForInput": True
}
})
mock_websocket.recv.return_value = raw_response_json
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
result = await session._receive()
# Assert the results
assert isinstance(result, types.LiveServerMessage)
assert result.server_content is not None
assert result.server_content.model_turn.parts[0].text == "Please provide more details."
assert result.server_content.turn_complete is True
assert result.server_content.turn_complete_reason == types.TurnCompleteReason.NEED_MORE_INPUT
assert result.server_content.waiting_for_input is True

View File

@@ -0,0 +1,147 @@
# 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.
#
"""Tests for live.py."""
import base64
import json
from unittest import mock
import pytest
from websockets import client
from .. import pytest_helper
from ... import client as gl_client
from ... import live
from ... import types
def mock_api_client(vertexai=False):
api_client = mock.MagicMock(spec=gl_client.BaseApiClient)
api_client.api_key = 'TEST_API_KEY'
api_client._host = lambda: 'test_host'
api_client._http_options = {'headers': {}} # Ensure headers exist
api_client.vertexai = vertexai
return api_client
@pytest.fixture
def mock_websocket():
websocket = mock.AsyncMock(spec=client.ClientConnection)
websocket.send = mock.AsyncMock()
websocket.recv = mock.AsyncMock(
return_value='{"serverContent": {"turnComplete": true}}'
) # Default response
websocket.close = mock.AsyncMock()
return websocket
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_content_dict(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
content = [{'parts': [{'text': 'test'}]}]
await session.send_client_content(turns=content)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'client_content' in sent_data
assert sent_data['client_content']['turns'][0]['parts'][0]['text'] == 'test'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_content_dict_list(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
content = [{'parts': [{'text': 'test'}]}]
await session.send_client_content(turns=content)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'client_content' in sent_data
assert sent_data['client_content']['turns'][0]['parts'][0]['text'] == 'test'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_content_content(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
content = types.Content.model_validate({'parts': [{'text': 'test'}]})
await session.send_client_content(turns=content)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'client_content' in sent_data
assert sent_data['client_content']['turns'][0]['parts'][0]['text'] == 'test'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_content_with_blob(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
content = types.Content.model_validate(
{'parts': [{'inline_data': {'data': b'test'}}]}
)
await session.send_client_content(turns=content)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'client_content' in sent_data
assert pytest_helper.get_value_ignore_key_case(
sent_data['client_content']['turns'][0]['parts'][0], 'inline_data') == {
'data': base64.b64encode(b'test').decode()
}
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_client_content_turn_complete_false(
mock_websocket, vertexai
):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
await session.send_client_content(turn_complete=False)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'client_content' in sent_data
assert sent_data['client_content']['turnComplete'] == False
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_client_content_empty(
mock_websocket, vertexai
):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
await session.send_client_content()
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'client_content' in sent_data
assert sent_data['client_content']['turnComplete'] == True

View File

@@ -0,0 +1,268 @@
# 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.
#
"""Tests for live.py."""
import json
import os
from unittest import mock
import pytest
from websockets import client
from ... import client as gl_client
from ... import live
from ... import types
from .. import pytest_helper
IMAGE_FILE_PATH = os.path.abspath(
os.path.join(os.path.dirname(__file__), '../data/google.jpg')
)
def mock_api_client(vertexai=False):
api_client = mock.MagicMock(spec=gl_client.BaseApiClient)
api_client.api_key = 'TEST_API_KEY'
api_client._host = lambda: 'test_host'
api_client._http_options = {'headers': {}} # Ensure headers exist
api_client.vertexai = vertexai
api_client._api_client = api_client
return api_client
@pytest.fixture
def mock_websocket():
websocket = mock.AsyncMock(spec=client.ClientConnection)
websocket.send = mock.AsyncMock()
websocket.recv = mock.AsyncMock(
return_value='{"serverContent": {"turnComplete": true}}'
) # Default response
websocket.close = mock.AsyncMock()
return websocket
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_media_blob_dict(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
content = {'data': bytes([0, 0, 0, 0, 0, 0]), 'mime_type': 'audio/pcm'}
await session.send_realtime_input(media=content)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert sent_data['realtime_input']['mediaChunks'][0]['data'] == 'AAAAAAAA'
assert pytest_helper.get_value_ignore_key_case(
sent_data['realtime_input']['mediaChunks'][0], 'mime_type'
) == 'audio/pcm'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_media_blob(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
content = types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='audio/pcm')
await session.send_realtime_input(media=content)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert sent_data['realtime_input']['mediaChunks'][0]['data'] == 'AAAAAAAA'
assert pytest_helper.get_value_ignore_key_case(
sent_data['realtime_input']['mediaChunks'][0], 'mime_type'
) == 'audio/pcm'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_media_image(mock_websocket, vertexai, image_jpeg):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
await session.send_realtime_input(media=image_jpeg)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert pytest_helper.get_value_ignore_key_case(
sent_data['realtime_input']['mediaChunks'][0], 'mime_type'
) == 'image/jpeg'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.parametrize(
'content',
[
{'data': bytes([0, 0, 0, 0, 0, 0]), 'mime_type': 'audio/pcm'},
types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='audio/pcm'),
],
)
@pytest.mark.asyncio
async def test_send_audio(mock_websocket, vertexai, content):
api_client = mock_api_client(vertexai=vertexai)
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
await session.send_realtime_input(audio=content)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert sent_data['realtime_input']['audio']['data'] == 'AAAAAAAA'
assert pytest_helper.get_value_ignore_key_case(
sent_data['realtime_input']['audio'], 'mime_type'
) == 'audio/pcm'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_bad_audio_blob(mock_websocket, vertexai):
api_client = mock_api_client(vertexai=vertexai)
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
content = types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='image/png')
with pytest.raises(ValueError, match='.*Unsupported mime type.*'):
await session.send_realtime_input(audio=content)
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_bad_video_blob(mock_websocket, vertexai):
api_client = mock_api_client(vertexai=vertexai)
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
content = types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='audio/pcm')
with pytest.raises(ValueError, match='.*Unsupported mime type.*'):
await session.send_realtime_input(video=content)
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_audio_stream_end(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
await session.send_realtime_input(audio_stream_end=True)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert sent_data['realtime_input']['audioStreamEnd'] == True
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.parametrize(
'content',
[
{'data': bytes([0, 0, 0, 0, 0, 0]), 'mime_type': 'image/png'},
types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='image/png'),
],
)
@pytest.mark.asyncio
async def test_send_video(mock_websocket, vertexai, content):
api_client = mock_api_client(vertexai=vertexai)
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
await session.send_realtime_input(video=content)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert sent_data['realtime_input']['video']['data'] == 'AAAAAAAA'
assert pytest_helper.get_value_ignore_key_case(
sent_data['realtime_input']['video'], 'mime_type'
) == 'image/png'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_video_image(mock_websocket, vertexai, image_jpeg):
api_client = mock_api_client(vertexai=vertexai)
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
await session.send_realtime_input(video=image_jpeg)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert pytest_helper.get_value_ignore_key_case(
sent_data['realtime_input']['video'], 'mime_type'
) == 'image/jpeg'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_text(mock_websocket, vertexai):
api_client = mock_api_client(vertexai=vertexai)
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
await session.send_realtime_input(text='Hello?')
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert sent_data['realtime_input']['text'] == 'Hello?'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.parametrize('activity', [{}, types.ActivityStart()])
@pytest.mark.asyncio
async def test_send_activity_start(mock_websocket, vertexai, activity):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
await session.send_realtime_input(activity_start=activity)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert sent_data['realtime_input']['activityStart'] == {}
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.parametrize('activity', [{}, types.ActivityEnd()])
@pytest.mark.asyncio
async def test_send_activity_end(mock_websocket, vertexai, activity):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
await session.send_realtime_input(activity_end=activity)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert sent_data['realtime_input']['activityEnd'] == {}
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_multiple_args(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
with pytest.raises(ValueError, match='.*one argument.*'):
await session.send_realtime_input(
text='Hello?', activity_start=types.ActivityStart()
)

View File

@@ -0,0 +1,222 @@
# 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.
#
"""Tests for live.py."""
import contextlib
import json
from unittest import mock
import pytest
from websockets import client
from .. import pytest_helper
from ... import client as gl_client
from ... import live
from ... import types
def exception_if_mldev(vertexai, exception_type: type[Exception]):
if vertexai:
return contextlib.nullcontext()
else:
return pytest.raises(exception_type)
def mock_api_client(vertexai=False):
api_client = mock.MagicMock(spec=gl_client.BaseApiClient)
api_client.api_key = 'TEST_API_KEY'
api_client._host = lambda: 'test_host'
api_client._http_options = {'headers': {}} # Ensure headers exist
api_client.vertexai = vertexai
api_client._api_client = api_client
return api_client
@pytest.fixture
def mock_websocket():
websocket = mock.AsyncMock(spec=client.ClientConnection)
websocket.send = mock.AsyncMock()
websocket.recv = mock.AsyncMock(
return_value='{"serverContent": {"turnComplete": true}}'
) # Default response
websocket.close = mock.AsyncMock()
return websocket
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_function_response_dict(mock_websocket, vertexai):
api_client = mock_api_client(vertexai=vertexai)
session = live.AsyncSession(
api_client=api_client, websocket=mock_websocket
)
input = {
'name': 'get_current_weather',
'response': {'temperature': 14.5, 'unit': 'C'},
}
if not vertexai:
input['id'] = 'some-id'
await session.send_tool_response(function_responses=input)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'tool_response' in sent_data
assert (
sent_data['tool_response']['functionResponses'][0]['name']
== 'get_current_weather'
)
assert (
sent_data['tool_response']['functionResponses'][0]['response'][
'temperature'
]
== 14.5
)
assert (
sent_data['tool_response']['functionResponses'][0]['response']['unit']
== 'C'
)
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_function_response(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
input = types.FunctionResponse(
name='get_current_weather',
response={
'temperature': 14.5,
'unit': 'C',
'user_name': 'test_user_name',
'userEmail': 'test_user_email',
},
)
if not vertexai:
input.id = 'some-id'
await session.send_tool_response(function_responses=input)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'tool_response' in sent_data
assert (
sent_data['tool_response']['functionResponses'][0]['name']
== 'get_current_weather'
)
assert (
sent_data['tool_response']['functionResponses'][0]['response']
== input.response
)
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_function_response_scheduling(mock_websocket, vertexai):
api_client = mock_api_client(vertexai=vertexai)
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
input = types.FunctionResponse(
name='get_current_weather',
response={'temperature': 14.5, 'unit': 'C'},
will_continue=True,
scheduling=types.FunctionResponseScheduling.SILENT,
)
if not vertexai:
input.id = 'some-id'
await session.send_tool_response(function_responses=input)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'tool_response' in sent_data
assert pytest_helper.get_value_ignore_key_case(
sent_data['tool_response']['functionResponses'][0], 'will_continue'
)
assert (
sent_data['tool_response']['functionResponses'][0]['scheduling']
== 'SILENT'
)
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_function_response_list(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
input1 = {
'name': 'get_current_weather',
'response': {'temperature': 14.5, 'unit': 'C'},
}
input2 = {
'name': 'get_current_weather',
'response': {'temperature': 99.9, 'unit': 'C'},
}
if not vertexai:
input1['id'] = '1'
input2['id'] = '2'
await session.send_tool_response(function_responses=[input1, input2])
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'tool_response' in sent_data
assert len(sent_data['tool_response']['functionResponses']) == 2
assert (
sent_data['tool_response']['functionResponses'][0]['response'][
'temperature'
]
== 14.5
)
assert (
sent_data['tool_response']['functionResponses'][1]['response'][
'temperature'
]
== 99.9
)
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_missing_id(mock_websocket, vertexai):
api_client = mock_api_client(vertexai=vertexai)
session = live.AsyncSession(
api_client=api_client, websocket=mock_websocket
)
input1 = {
'name': 'get_current_weather',
'response': {'temperature': 14.5, 'unit': 'C'},
'id': '1',
}
input2 = {
'name': 'get_current_weather',
'response': {'temperature': 99.9, 'unit': 'C'},
}
if not vertexai:
with pytest.raises(ValueError, match=".*must have.*"):
await session.send_tool_response(function_responses=[input1, input2])

View File

@@ -0,0 +1,17 @@
# 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.
#
"""Tests for the Google GenAI SDK's local_tokenizer module."""

View File

@@ -0,0 +1,343 @@
# 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 unittest
from unittest.mock import MagicMock, patch
from sentencepiece import sentencepiece_model_pb2
from ... import local_tokenizer
from ... import types
class TestLocalTokenizer(unittest.TestCase):
def setUp(self):
# This setup will be used by all tests
self.mock_load_model_proto = patch(
'genai._local_tokenizer_loader.load_model_proto'
).start()
self.mock_get_sentencepiece = patch(
'genai._local_tokenizer_loader.get_sentencepiece'
).start()
self.mock_load_model_proto.return_value = MagicMock()
self.mock_tokenizer = MagicMock()
self.mock_get_sentencepiece.return_value = self.mock_tokenizer
self.tokenizer = local_tokenizer.LocalTokenizer(model_name='gemini-3-pro-preview')
def tearDown(self):
patch.stopall()
def test_count_tokens_simple_string(self):
self.mock_tokenizer.encode.return_value = [[1, 2, 3]]
result = self.tokenizer.count_tokens('Hello world')
self.assertEqual(result.total_tokens, 3)
self.mock_tokenizer.encode.assert_called_once_with(['Hello world'])
def test_count_tokens_list_of_strings(self):
self.mock_tokenizer.encode.return_value = [[1, 2], [3]]
result = self.tokenizer.count_tokens(['Hello', 'world'])
self.assertEqual(result.total_tokens, 3)
self.mock_tokenizer.encode.assert_called_once_with(['Hello', 'world'])
def test_count_tokens_with_content_object(self):
self.mock_tokenizer.encode.return_value = [[1, 2, 3]]
content = types.Content(parts=[types.Part(text='Hello world')])
result = self.tokenizer.count_tokens(content)
self.assertEqual(result.total_tokens, 3)
self.mock_tokenizer.encode.assert_called_once_with(['Hello world'])
def test_count_tokens_with_chat_history(self):
self.mock_tokenizer.encode.return_value = [[1, 2], [3, 4, 5]]
history = [
types.Content(role='user', parts=[types.Part(text='Hello')]),
types.Content(role='model', parts=[types.Part(text='Hi there!')]),
]
result = self.tokenizer.count_tokens(history)
self.assertEqual(result.total_tokens, 5)
self.mock_tokenizer.encode.assert_called_once_with(['Hello', 'Hi there!'])
def test_count_tokens_with_tools(self):
self.mock_tokenizer.encode.return_value = [
[1],
[1, 2],
[1, 2, 3],
[1, 2, 3, 4],
[1, 2, 3, 4, 5],
[1, 2, 3, 4, 5, 6],
]
tool = types.Tool(
function_declarations=[
types.FunctionDeclaration(
name='get_weather',
description='Get the weather for a location',
parameters=types.Schema(
type=types.Type.OBJECT,
properties={
'location': types.Schema(
type=types.Type.STRING, description='The location'
)
},
required=['location'],
),
)
]
)
config = types.CountTokensConfig(tools=[tool])
result = self.tokenizer.count_tokens(
'What is the weather in Boston?', config=config
)
self.assertEqual(result.total_tokens, 21)
self.mock_tokenizer.encode.assert_called_once_with([
'What is the weather in Boston?',
'get_weather',
'Get the weather for a location',
'location',
'location',
'The location',
])
def test_count_tokens_with_function_call(self):
self.mock_tokenizer.encode.return_value = [[1, 2], [3], [4, 5]]
content = types.Content(
role='model',
parts=[
types.Part(
function_call=types.FunctionCall(
name='get_weather', args={'location': 'Boston'}
)
)
],
)
result = self.tokenizer.count_tokens(content)
self.assertEqual(result.total_tokens, 5)
self.mock_tokenizer.encode.assert_called_once_with(
['get_weather', 'location', 'Boston']
)
def test_count_tokens_with_function_response(self):
self.mock_tokenizer.encode.return_value = [[1, 2], [3], [4, 5]]
content = types.Content(
role='user',
parts=[
types.Part(
function_response=types.FunctionResponse(
name='get_weather', response={'weather': 'sunny'}
)
)
],
)
result = self.tokenizer.count_tokens(content)
self.assertEqual(result.total_tokens, 5)
self.mock_tokenizer.encode.assert_called_once_with(
['get_weather', 'weather', 'sunny']
)
def test_count_tokens_with_unsupported_content(self):
with self.assertRaises(ValueError):
self.tokenizer.count_tokens(
[
types.Content(
parts=[
types.Part(
inline_data=types.Blob(
data=b'test', mime_type='image/png'
)
)
]
)
]
)
def test_count_tokens_with_system_instruction(self):
self.mock_tokenizer.encode.return_value = [[1, 2, 3], [4, 5]]
config = types.CountTokensConfig(
system_instruction=types.Content(
parts=[types.Part(text='You are a helpful assistant.')]
)
)
result = self.tokenizer.count_tokens('Hello', config=config)
self.assertEqual(result.total_tokens, 5)
self.mock_tokenizer.encode.assert_called_once_with(
['Hello', 'You are a helpful assistant.']
)
def test_count_tokens_with_response_schema(self):
self.mock_tokenizer.encode.return_value = [
[1],
[1, 2],
[1, 2, 3],
[1, 2, 3, 4],
[1, 2, 3, 4, 5],
]
schema = types.Schema(
type=types.Type.OBJECT,
format='schema_format',
description='Recipe schema',
enum=['schema_enum1', 'schema_enum2'],
properties={
'recipe_name': types.Schema(
type=types.Type.STRING,
description='Name of the recipe',
)
},
items=types.Schema(
type=types.Type.STRING,
description='Item in the recipe',
),
example={
'recipe_example': types.Schema(
type=types.Type.STRING,
description='example in the recipe',
)
},
required=['recipe_name'],
)
config = types.CountTokensConfig(
generation_config=types.GenerationConfig(response_schema=schema)
)
result = self.tokenizer.count_tokens(
'Generate a recipe for chocolate chip cookies.', config=config
)
self.assertEqual(result.total_tokens, 15)
self.mock_tokenizer.encode.assert_called_once_with([
'Generate a recipe for chocolate chip cookies.',
'schema_format',
'Recipe schema',
'schema_enum1',
'schema_enum2',
'recipe_name',
'Item in the recipe',
'recipe_name',
'Name of the recipe',
'recipe_example',
])
def test_count_tokens_with_unsupported_fields_logs_warning(self):
self.mock_tokenizer.encode.return_value = [[1, 2, 3]]
content_with_unsupported = types.Content(
role='user',
parts=[
types.Part(text='hello'),
# executable_code is not supported by _TextsAccumulator
types.Part(
executable_code=types.ExecutableCode(
language='PYTHON', code='print(1)'
)
),
],
)
with self.assertLogs('google_genai.local_tokenizer', level='WARNING') as cm:
self.tokenizer.count_tokens(content_with_unsupported)
self.assertIn(
'Content contains unsupported types for token counting', cm.output[0]
)
def test_compute_tokens_simple_string(self):
mock_spt = MagicMock()
mock_spt.pieces = [
MagicMock(id=1, piece='He'),
MagicMock(id=2, piece='llo'),
MagicMock(id=3, piece=' world'),
]
self.mock_tokenizer.EncodeAsImmutableProto.return_value = [mock_spt]
result = self.tokenizer.compute_tokens('Hello world')
self.assertEqual(len(result.tokens_info), 1)
self.assertEqual(result.tokens_info[0].token_ids, [1, 2, 3])
self.assertEqual(result.tokens_info[0].tokens, [b'He', b'llo', b' world'])
self.assertEqual(result.tokens_info[0].role, 'user')
self.mock_tokenizer.EncodeAsImmutableProto.assert_called_once_with(
['Hello world']
)
def test_compute_tokens_with_chat_history(self):
mock_spt1 = MagicMock()
mock_spt1.pieces = [MagicMock(id=1, piece='Hello')]
mock_spt2 = MagicMock()
mock_spt2.pieces = [
MagicMock(id=2, piece='Hi'),
MagicMock(id=3, piece=' there!'),
]
self.mock_tokenizer.EncodeAsImmutableProto.return_value = [
mock_spt1,
mock_spt2,
]
history = [
types.Content(role='user', parts=[types.Part(text='Hello')]),
types.Content(role='model', parts=[types.Part(text='Hi there!')]),
]
result = self.tokenizer.compute_tokens(history)
self.assertEqual(len(result.tokens_info), 2)
self.assertEqual(result.tokens_info[0].token_ids, [1])
self.assertEqual(result.tokens_info[0].tokens, [b'Hello'])
self.assertEqual(result.tokens_info[0].role, 'user')
self.assertEqual(result.tokens_info[1].token_ids, [2, 3])
self.assertEqual(result.tokens_info[1].tokens, [b'Hi', b' there!'])
self.assertEqual(result.tokens_info[1].role, 'model')
self.mock_tokenizer.EncodeAsImmutableProto.assert_called_once_with(
['Hello', 'Hi there!']
)
def test_compute_tokens_with_byte_tokens(self):
mock_spt = MagicMock()
mock_spt.pieces = [
MagicMock(id=1, piece='<0x48>'),
MagicMock(id=2, piece='ello'),
]
self.mock_tokenizer.EncodeAsImmutableProto.return_value = [mock_spt]
self.tokenizer._model_proto = sentencepiece_model_pb2.ModelProto(
pieces=[
sentencepiece_model_pb2.ModelProto.SentencePiece(),
sentencepiece_model_pb2.ModelProto.SentencePiece(
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.BYTE
),
sentencepiece_model_pb2.ModelProto.SentencePiece(
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.NORMAL
),
]
)
result = self.tokenizer.compute_tokens('Hello')
self.assertEqual(len(result.tokens_info), 1)
self.assertEqual(result.tokens_info[0].token_ids, [1, 2])
self.assertEqual(result.tokens_info[0].tokens, [b'H', b'ello'])
self.mock_tokenizer.EncodeAsImmutableProto.assert_called_once_with(
['Hello']
)
class TestParseHexByte(unittest.TestCase):
def test_valid_hex(self):
self.assertEqual(local_tokenizer._parse_hex_byte('<0x41>'), 65)
self.assertEqual(local_tokenizer._parse_hex_byte('<0xFF>'), 255)
self.assertEqual(local_tokenizer._parse_hex_byte('<0x00>'), 0)
def test_invalid_length(self):
with self.assertRaisesRegex(ValueError, 'Invalid byte length'):
local_tokenizer._parse_hex_byte('<0x41')
with self.assertRaisesRegex(ValueError, 'Invalid byte length'):
local_tokenizer._parse_hex_byte('<0x411>')
def test_invalid_format(self):
with self.assertRaisesRegex(ValueError, 'Invalid byte format'):
local_tokenizer._parse_hex_byte(' 0x41>')
with self.assertRaisesRegex(ValueError, 'Invalid byte format'):
local_tokenizer._parse_hex_byte('<0x41 ')
def test_invalid_hex_value(self):
with self.assertRaisesRegex(ValueError, 'Invalid hex value'):
local_tokenizer._parse_hex_byte('<0xFG>')

View File

@@ -0,0 +1,235 @@
# 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 unittest
from unittest.mock import MagicMock, mock_open, patch
import sentencepiece as spm
from sentencepiece import sentencepiece_model_pb2
from ... import _local_tokenizer_loader as loader
# A minimal valid sentencepiece model proto
FAKE_MODEL_CONTENT = sentencepiece_model_pb2.ModelProto(
pieces=[
sentencepiece_model_pb2.ModelProto.SentencePiece(
piece="<unk>",
score=0,
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.UNKNOWN,
),
sentencepiece_model_pb2.ModelProto.SentencePiece(
piece="<s>",
score=0,
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.CONTROL,
),
sentencepiece_model_pb2.ModelProto.SentencePiece(
piece="</s>",
score=0,
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.CONTROL,
),
sentencepiece_model_pb2.ModelProto.SentencePiece(
piece="a",
score=0,
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.NORMAL,
),
]
).SerializeToString()
GEMMA2_HASH = "61a7b147390c64585d6c3543dd6fc636906c9af3865a5548f27f31aee1d4c8e2"
class TestGetTokenizerName(unittest.TestCase):
def test_get_tokenizer_name_success(self):
self.assertEqual(loader.get_tokenizer_name("gemini-2.5-pro"), "gemma3")
self.assertEqual(
loader.get_tokenizer_name("gemini-2.5-pro-preview-06-05"), "gemma3"
)
def test_get_tokenizer_name_unsupported(self):
with self.assertRaisesRegex(
ValueError, "Model unsupported-model is not supported"
):
loader.get_tokenizer_name("unsupported-model")
@patch("genai._local_tokenizer_loader.os.rename")
@patch("genai._local_tokenizer_loader.os.makedirs")
@patch("genai._local_tokenizer_loader.os.remove")
@patch("genai._local_tokenizer_loader.open", new_callable=mock_open)
@patch("genai._local_tokenizer_loader.os.path.exists")
@patch("genai._local_tokenizer_loader.requests.get")
@patch("genai._local_tokenizer_loader.hashlib.sha256")
class TestLoaderFunctions(unittest.TestCase):
def setUp(self):
# Clear caches before each test
loader.load_model_proto.cache_clear()
loader.get_sentencepiece.cache_clear()
# Patch tempfile.gettempdir to control cache location
self.tempdir_patcher = patch(
"tempfile.gettempdir", return_value="/tmp/fake_temp_dir"
)
self.mock_tempdir = self.tempdir_patcher.start()
def tearDown(self):
self.tempdir_patcher.stop()
def _setup_get_mock(self, mock_get):
mock_response = MagicMock()
mock_response.content = FAKE_MODEL_CONTENT
mock_response.raise_for_status.return_value = None
mock_get.return_value = mock_response
def test_load_model_proto_from_url(
self,
mock_sha256,
mock_get,
mock_exists,
mock_open_func,
mock_remove,
mock_makedirs,
mock_rename,
):
mock_exists.return_value = False # Don't use cache
self._setup_get_mock(mock_get)
mock_sha256.return_value.hexdigest.return_value = GEMMA2_HASH
proto = loader.load_model_proto("gemma2")
self.assertIsInstance(proto, sentencepiece_model_pb2.ModelProto)
self.assertEqual(len(proto.pieces), 4)
mock_get.assert_called_once()
mock_makedirs.assert_called_once()
mock_open_func.assert_called()
mock_rename.assert_called_once()
def test_load_model_proto_from_cache(
self,
mock_sha256,
mock_get,
mock_exists,
mock_open_func,
mock_remove,
mock_makedirs,
mock_rename,
):
mock_exists.return_value = True # Use cache
mock_open_func.return_value.read.return_value = FAKE_MODEL_CONTENT
mock_sha256.return_value.hexdigest.return_value = GEMMA2_HASH
proto = loader.load_model_proto("gemma2")
self.assertIsInstance(proto, sentencepiece_model_pb2.ModelProto)
mock_get.assert_not_called()
def test_load_model_proto_corrupted_cache(
self,
mock_sha256,
mock_get,
mock_exists,
mock_open_func,
mock_remove,
mock_makedirs,
mock_rename,
):
mock_exists.return_value = True # Use cache initially
self._setup_get_mock(mock_get)
mock_open_func.return_value.__enter__.return_value.read.return_value = (
b"corrupted"
)
# First hash for corrupted cache, second for good download
mock_sha256.side_effect = [
MagicMock(hexdigest=MagicMock(return_value="wrong_hash")),
MagicMock(hexdigest=MagicMock(return_value=GEMMA2_HASH)),
]
proto = loader.load_model_proto("gemma2")
self.assertIsInstance(proto, sentencepiece_model_pb2.ModelProto)
mock_remove.assert_called_once()
mock_get.assert_called_once()
def test_load_model_proto_bad_hash_from_url(
self,
mock_sha256,
mock_get,
mock_exists,
mock_open_func,
mock_remove,
mock_makedirs,
mock_rename,
):
mock_exists.return_value = False
self._setup_get_mock(mock_get)
mock_sha256.return_value.hexdigest.return_value = "wrong_hash"
with self.assertRaisesRegex(
ValueError, "Downloaded model file is corrupted"
):
loader.load_model_proto("gemma2")
def test_load_model_proto_unsupported(self, *args):
with self.assertRaisesRegex(
ValueError, "Tokenizer unsupported is not supported"
):
loader.load_model_proto("unsupported")
def test_get_sentencepiece_success(
self,
mock_sha256,
mock_get,
mock_exists,
mock_open_func,
mock_remove,
mock_makedirs,
mock_rename,
):
mock_exists.return_value = False
self._setup_get_mock(mock_get)
mock_sha256.return_value.hexdigest.return_value = GEMMA2_HASH
processor = loader.get_sentencepiece("gemma2")
self.assertIsInstance(processor, spm.SentencePieceProcessor)
mock_get.assert_called_once()
def test_get_sentencepiece_unsupported(self, *args):
with self.assertRaisesRegex(
ValueError, "Tokenizer unsupported is not supported"
):
loader.get_sentencepiece("unsupported")
def test_get_sentencepiece_caching(
self,
mock_sha256,
mock_get,
mock_exists,
mock_open_func,
mock_remove,
mock_makedirs,
mock_rename,
):
mock_exists.return_value = False
self._setup_get_mock(mock_get)
mock_sha256.return_value.hexdigest.return_value = GEMMA2_HASH
# Call twice
loader.get_sentencepiece("gemma2")
loader.get_sentencepiece("gemma2")
# Should only be loaded once due to lru_cache
mock_get.assert_called_once()

View File

@@ -0,0 +1,17 @@
# 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.
#
"""Tests for the Google GenAI SDK's _mcp_utils module."""

View File

@@ -0,0 +1,89 @@
# 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 typing
from typing import Any
from ... import _mcp_utils
from ... import types
if typing.TYPE_CHECKING:
from mcp.types import Tool as McpTool
from mcp import ClientSession as McpClientSession
else:
McpTool: typing.Type = Any
McpClientSession: typing.Type = Any
try:
from mcp.types import Tool as McpTool
from mcp import ClientSession as McpClientSession
except ImportError:
McpTool = None
McpClientSession = None
def test_mcp_tools():
"""Test whether the list of tools contains any MCP tools."""
if McpTool is None:
return
mcp_tools = [
McpTool(
name='tool',
description='tool-description',
inputSchema={
'type': 'OBJECT',
'properties': {
'key1': {'type': 'STRING'},
'key2': {'type': 'NUMBER'},
},
},
),
]
assert _mcp_utils.has_mcp_tool_usage(mcp_tools)
def test_mcp_client_session():
"""Test whether the list of tools contains any MCP tools."""
class MockMcpClientSession(McpClientSession):
def __init__(self):
self._read_stream = None
self._write_stream = None
mcp_tools = [
MockMcpClientSession(),
]
assert _mcp_utils.has_mcp_tool_usage(mcp_tools)
def test_no_mcp_tools():
if McpClientSession is None:
return
"""Test whether the list of tools contains any MCP tools."""
gemini_tools = [
types.Tool(
function_declarations=[
types.FunctionDeclaration(
name='tool',
description='tool-description',
parameters=types.Schema(
type='OBJECT',
properties={},
),
),
],
),
]
assert not _mcp_utils.has_mcp_tool_usage(gemini_tools)

View File

@@ -0,0 +1,191 @@
# 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.
#
from ... import _mcp_utils
from ... import types
try:
from mcp import types as mcp_types
except ImportError as e:
import sys
if sys.version_info < (3, 10):
raise ImportError(
'MCP Tool requires Python 3.10 or above. Please upgrade your Python'
' version.'
) from e
else:
raise e
def test_empty_mcp_tools_list():
"""Test conversion of empty MCP tools list to Gemini tools list."""
result = _mcp_utils.mcp_to_gemini_tools([])
assert result == []
def test_unknown_field_conversion():
"""Test conversion of MCP tools with unknown fields to Gemini tools."""
mcp_tools = [
mcp_types.Tool(
name='tool',
description='tool-description',
inputSchema={
'type': 'object',
'properties': {},
'unknown_field': 'unknownField',
'unknown_object': {},
},
),
]
result = _mcp_utils.mcp_to_gemini_tools(mcp_tools)
assert result == [
types.Tool(
function_declarations=[
types.FunctionDeclaration(
name='tool',
description='tool-description',
parameters=types.Schema(
type='OBJECT',
properties={},
),
),
],
),
]
def test_items_conversion():
"""Test conversion of MCP tools with items to Gemini tools."""
mcp_tools = [
mcp_types.Tool(
name='tool',
description='tool-description',
inputSchema={
'type': 'array',
'items': {
'type': 'object',
'properties': {
'key1': {
'type': 'string',
},
'key2': {
'type': 'number',
},
},
},
},
),
]
result = _mcp_utils.mcp_to_gemini_tools(mcp_tools)
assert result == [
types.Tool(
function_declarations=[
types.FunctionDeclaration(
name='tool',
description='tool-description',
parameters=types.Schema(
type='ARRAY',
items=types.Schema(
type='OBJECT',
properties={
'key1': types.Schema(type='STRING'),
'key2': types.Schema(type='NUMBER'),
},
),
),
),
],
),
]
def test_any_of_conversion():
"""Test conversion of MCP tools with any_of to Gemini tools."""
mcp_tools = [
mcp_types.Tool(
name='tool',
description='tool-description',
inputSchema={
'type': 'object',
'any_of': [
{
'type': 'string',
},
{
'type': 'number',
},
],
},
),
]
result = _mcp_utils.mcp_to_gemini_tools(mcp_tools)
assert result == [
types.Tool(
function_declarations=[
types.FunctionDeclaration(
name='tool',
description='tool-description',
parameters=types.Schema(
type='OBJECT',
any_of=[
types.Schema(type='STRING'),
types.Schema(type='NUMBER'),
],
),
),
],
),
]
def test_properties_conversion():
"""Test conversion of MCP tools with properties to Gemini tools."""
mcp_tools = [
mcp_types.Tool(
name='tool',
description='tool-description',
inputSchema={
'type': 'object',
'properties': {
'key1': {
'type': 'string',
},
'key2': {
'type': 'number',
},
},
},
),
]
result = _mcp_utils.mcp_to_gemini_tools(mcp_tools)
assert result == [
types.Tool(
function_declarations=[
types.FunctionDeclaration(
name='tool',
description='tool-description',
parameters=types.Schema(
type='OBJECT',
properties={
'key1': types.Schema(type='STRING'),
'key2': types.Schema(type='NUMBER'),
},
),
),
],
),
]

View File

@@ -0,0 +1,201 @@
# 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 _asyncio
import pytest
from ... import _extra_utils
from ... import types
from ..._adapters import McpToGenAiToolAdapter
try:
from mcp import types as mcp_types
from mcp import ClientSession as McpClientSession
except ImportError as e:
import sys
if sys.version_info < (3, 10):
raise ImportError(
'MCP Tool requires Python 3.10 or above. Please upgrade your Python'
' version.'
) from e
else:
raise e
@pytest.mark.asyncio
async def test_parse_empty_config_dict():
"""Test conversion of empty GenerateContentConfigDict to parsed config."""
config = {}
parsed_config, mcp_to_genai_tool_adapters = (
await _extra_utils.parse_config_for_mcp_sessions(config)
)
assert parsed_config == None
assert not mcp_to_genai_tool_adapters
@pytest.mark.asyncio
async def test_parse_empty_config_object():
"""Test conversion of empty GenerateContentConfig to parsed config."""
config = types.GenerateContentConfig()
parsed_config, mcp_to_genai_tool_adapters = (
await _extra_utils.parse_config_for_mcp_sessions(config)
)
assert config is not parsed_config # config is not modified
assert not mcp_to_genai_tool_adapters
@pytest.mark.asyncio
async def test_parse_config_object_with_tools():
"""Test conversion of GenerateContentConfig with tools to parsed config."""
class MockMcpClientSession(McpClientSession):
def __init__(self):
self._read_stream = None
self._write_stream = None
async def list_tools(self):
return mcp_types.ListToolsResult(
tools=[
mcp_types.Tool(
name='get_weather',
description='Get the weather in a city.',
inputSchema={
'type': 'object',
'properties': {'location': {'type': 'string'}},
},
),
mcp_types.Tool(
name='get_weather_2',
description='Different tool to get the weather.',
inputSchema={
'type': 'object',
'properties': {'location': {'type': 'string'}},
},
),
]
)
mock_session_instance = MockMcpClientSession()
config = types.GenerateContentConfig(tools=[mock_session_instance])
parsed_config, mcp_to_genai_tool_adapters = (
await _extra_utils.parse_config_for_mcp_sessions(config)
)
assert len(config.tools) == 1
assert config.tools[0] is mock_session_instance
assert config is not parsed_config # config is not modified
assert len(mcp_to_genai_tool_adapters) == 2
assert mcp_to_genai_tool_adapters.keys() == {
'get_weather',
'get_weather_2',
}
assert isinstance(
mcp_to_genai_tool_adapters['get_weather'], McpToGenAiToolAdapter
)
assert isinstance(
mcp_to_genai_tool_adapters['get_weather_2'], McpToGenAiToolAdapter
)
@pytest.mark.asyncio
async def test_parse_config_object_with_tools_complex_type():
"""Test conversion of GenerateContentConfig with tools to parsed config."""
class MockMcpClientSession(McpClientSession):
def __init__(self):
self._read_stream = None
self._write_stream = None
self._future = _asyncio.Future() # This object cannot be pickled.
async def list_tools(self):
return mcp_types.ListToolsResult(
tools=[
mcp_types.Tool(
name='get_weather',
description='Get the weather in a city.',
inputSchema={
'type': 'object',
'properties': {'location': {'type': 'string'}},
},
),
mcp_types.Tool(
name='get_weather_2',
description='Different tool to get the weather.',
inputSchema={
'type': 'object',
'properties': {'location': {'type': 'string'}},
},
),
]
)
mock_session_instance = MockMcpClientSession()
config = types.GenerateContentConfig(tools=[mock_session_instance])
parsed_config, mcp_to_genai_tool_adapters = (
await _extra_utils.parse_config_for_mcp_sessions(config)
)
assert len(config.tools) == 1
assert config.tools[0] is mock_session_instance
assert config is not parsed_config # config is not modified
assert len(mcp_to_genai_tool_adapters) == 2
assert mcp_to_genai_tool_adapters.keys() == {
'get_weather',
'get_weather_2',
}
assert isinstance(
mcp_to_genai_tool_adapters['get_weather'], McpToGenAiToolAdapter
)
assert isinstance(
mcp_to_genai_tool_adapters['get_weather_2'], McpToGenAiToolAdapter
)
@pytest.mark.asyncio
async def test_parse_config_object_with_non_mcp_tools():
"""Test conversion of GenerateContentConfig with regular tools to parsed config."""
config = types.GenerateContentConfig(
tools=[
types.Tool(
function_declarations=[
{'name': 'tool-1', 'description': 'tool-1-description'}
]
),
types.Tool(
function_declarations=[
{'name': 'tool-2', 'description': 'tool-2-description'}
]
),
]
)
parsed_config, mcp_to_genai_tool_adapters = (
await _extra_utils.parse_config_for_mcp_sessions(config)
)
assert len(config.tools) == 2
assert config is not parsed_config # config is not modified
assert mcp_to_genai_tool_adapters == {}
assert parsed_config.tools == [
types.Tool(
function_declarations=[
{'name': 'tool-1', 'description': 'tool-1-description'}
]
),
types.Tool(
function_declarations=[
{'name': 'tool-2', 'description': 'tool-2-description'}
]
),
]

View File

@@ -0,0 +1,130 @@
# 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 re
import pytest
from ... import _extra_utils
from ... import types
from ..._adapters import McpToGenAiToolAdapter
try:
from mcp import types as mcp_types
from mcp import ClientSession as McpClientSession
except ImportError as e:
import sys
if sys.version_info < (3, 10):
raise ImportError(
'MCP Tool requires Python 3.10 or above. Please upgrade your Python'
' version.'
) from e
else:
raise e
def test_parse_empty_config_dict():
"""Test conversion of empty GenerateContentConfigDict to parsed config."""
config = {}
parsed_config = _extra_utils.parse_config_for_mcp_usage(config)
assert parsed_config == None
def test_parse_empty_config_object():
"""Test conversion of empty GenerateContentConfig to parsed config."""
config = types.GenerateContentConfig()
parsed_config = _extra_utils.parse_config_for_mcp_usage(config)
assert config is not parsed_config # config is not modified
def test_parse_config_object_with_tools():
"""Test conversion of GenerateContentConfig with tools to parsed config."""
class MockMcpClientSession(McpClientSession):
def __init__(self):
self._read_stream = None
self._write_stream = None
async def list_tools(self):
return mcp_types.ListToolsResult(
tools=[
mcp_types.Tool(
name='get_weather',
description='Get the weather in a city.',
inputSchema={
'type': 'object',
'properties': {'location': {'type': 'string'}},
},
),
mcp_types.Tool(
name='get_weather_2',
description='Different tool to get the weather.',
inputSchema={
'type': 'object',
'properties': {'location': {'type': 'string'}},
},
),
]
)
mock_session_instance = MockMcpClientSession()
config = types.GenerateContentConfig(tools=[mock_session_instance])
parsed_config = _extra_utils.parse_config_for_mcp_usage(config)
assert config.http_options is None
assert config is not parsed_config # config is not modified
assert re.match(
r'mcp_used/\d+\.\d+\.\d+',
parsed_config.http_options.headers['x-goog-api-client'],
)
def test_parse_config_object_with_tools_and_existing_headers():
"""Test conversion of GenerateContentConfig with tools and existing headers to parsed config."""
config = types.GenerateContentConfig(
tools=[
mcp_types.Tool(
name='get_weather',
description='Get the weather in a city.',
inputSchema={
'type': 'object',
'properties': {'location': {'type': 'string'}},
},
),
mcp_types.Tool(
name='get_weather_2',
description='Different tool to get the weather.',
inputSchema={
'type': 'object',
'properties': {'location': {'type': 'string'}},
},
),
],
http_options=types.HttpOptions(
headers={
'x-goog-api-client': 'google-genai-sdk/1.0.0 gl-python/1.0.0'
}
),
)
parsed_config = _extra_utils.parse_config_for_mcp_usage(config)
assert not re.match(
r'mcp_used/\d+\.\d+\.\d+',
config.http_options.headers['x-goog-api-client'],
)
assert config is not parsed_config # config is not modified
assert re.match(
r'google-genai-sdk/1.0.0 gl-python/1.0.0 mcp_used/\d+\.\d+\.\d+',
parsed_config.http_options.headers['x-goog-api-client'],
)

View File

@@ -0,0 +1,72 @@
# 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.
#
from importlib.metadata import version
import re
import typing
from typing import Any
from ... import _mcp_utils
from ... import types
_is_mcp_imported = False
if typing.TYPE_CHECKING:
import mcp
_is_mcp_imported = True
else:
try:
import mcp
_is_mcp_imported = True
except ImportError:
_is_mcp_imported = False
def test_set_mcp_usage_header_from_empty_dict():
if not _is_mcp_imported:
return
"""Test whether the MCP usage header is set correctly from an empty dict."""
headers = {}
_mcp_utils.set_mcp_usage_header(headers)
assert re.match(r'mcp_used/\d+\.\d+\.\d+', headers['x-goog-api-client'])
def test_set_mcp_usage_header_with_existing_header():
if not _is_mcp_imported:
return
"""Test whether the MCP usage header is set correctly from an existing header."""
headers = {'x-goog-api-client': 'google-genai-sdk/1.0.0 gl-python/1.0.0'}
_mcp_utils.set_mcp_usage_header(headers)
assert re.match(
r'google-genai-sdk/1.0.0 gl-python/1.0.0 mcp_used/\d+\.\d+\.\d+',
headers['x-goog-api-client'],
)
def test_set_mcp_usage_header_with_existing_mcp_header():
if not _is_mcp_imported:
return
"""Test whether the MCP usage header is set correctly from an existing MCP header."""
headers = {
'x-goog-api-client': (
'google-genai-sdk/1.0.0 gl-python/1.0.0 mcp_used/1.0.0'
)
}
_mcp_utils.set_mcp_usage_header(headers)
assert re.match(
r'google-genai-sdk/1.0.0 gl-python/1.0.0 mcp_used/\d+\.\d+\.\d+',
headers['x-goog-api-client'],
)

View File

@@ -0,0 +1,17 @@
# 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.
#
"""Tests for the Google GenAI SDK's models module."""

View File

@@ -0,0 +1,8 @@
VERTEX_HTTP_OPTIONS = {
'api_version': 'v1beta1',
'base_url': 'https://us-central1-aiplatform.googleapis.com/',
}
MLDEV_HTTP_OPTIONS = {
'api_version': 'v1beta',
'base_url': 'https://generativelanguage.googleapis.com/',
}

View File

@@ -0,0 +1,120 @@
# 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 pytest
from ... import _transformers as t
from ... import types
from .. import pytest_helper
from . import constants
_COMPUTE_TOKENS_PARAMS = types._ComputeTokensParameters(
model='gemini-2.5-flash',
contents=[t.t_content('Tell me a story in 300 words.')],
)
_COMPUTE_TOKENS_PARAMS_VERTEX_CUSTOM_URL = types._ComputeTokensParameters(
model='gemini-2.5-flash',
contents=[t.t_content('Tell me a story in 300 words.')],
config={'http_options': constants.VERTEX_HTTP_OPTIONS},
)
_COMPUTE_TOKENS_PARAMS_MLDEV_CUSTOM_URL = types._ComputeTokensParameters(
model='gemini-2.5-flash',
contents=[t.t_content('Tell me a story in 300 words.')],
config={'http_options': constants.MLDEV_HTTP_OPTIONS},
)
_UNICODE_STRING = '这是一条unicode测试🤪❤★'
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_compute_tokens',
exception_if_mldev='only supported in',
parameters=types._ComputeTokensParameters(
model='gemini-2.5-flash',
contents=[t.t_content('Tell me a story in 300 words.')],
),
),
pytest_helper.TestTableItem(
name='test_compute_tokens_vertex_custom_url',
parameters=_COMPUTE_TOKENS_PARAMS_VERTEX_CUSTOM_URL,
exception_if_mldev='only supported in',
),
pytest_helper.TestTableItem(
name='test_compute_tokens_mldev_custom_url',
parameters=_COMPUTE_TOKENS_PARAMS_MLDEV_CUSTOM_URL,
exception_if_vertex='404',
exception_if_mldev='only supported in',
),
pytest_helper.TestTableItem(
name='test_compute_tokens_unicode',
exception_if_mldev='only supported in',
parameters=types._ComputeTokensParameters(
model='gemini-2.5-flash', contents=[t.t_content(_UNICODE_STRING)]
),
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='models.compute_tokens',
test_table=test_table,
)
def test_token_bytes_deserialization(client):
if client._api_client.vertexai:
response = client.models.compute_tokens(
model=_COMPUTE_TOKENS_PARAMS.model,
contents=_UNICODE_STRING,
)
decoded_tokens = b''.join(response.tokens_info[0].tokens)
assert (
decoded_tokens
== b'\xe8\xbf\x99\xe6\x98\xaf\xe4\xb8\x80\xe6\x9d\xa1unicode\xe6\xb5\x8b\xe8\xaf\x95\xf0\x9f\xa4\xaa\xe2\x9d\xa4\xe2\x98\x85'
)
assert decoded_tokens.decode('utf-8') == _UNICODE_STRING
@pytest.mark.asyncio
async def test_async(client):
if client._api_client.vertexai:
response = await client.aio.models.compute_tokens(
model=_COMPUTE_TOKENS_PARAMS.model,
contents=_COMPUTE_TOKENS_PARAMS.contents,
)
assert response
else:
with pytest.raises(Exception):
await client.aio.models.compute_tokens(
model=_COMPUTE_TOKENS_PARAMS.model,
contents=_COMPUTE_TOKENS_PARAMS.contents,
)
def test_different_model_names(client):
if client._api_client.vertexai:
response1 = client.models.compute_tokens(
model='gemini-2.5-flash', contents=_COMPUTE_TOKENS_PARAMS.contents
)
assert response1
response3 = client.models.compute_tokens(
model='publishers/google/models/gemini-2.5-flash',
contents=_COMPUTE_TOKENS_PARAMS.contents,
)
assert response3
response4 = client.models.compute_tokens(
model='projects/vertexsdk/locations/us-central1/publishers/google/models/gemini-2.5-flash',
contents=_COMPUTE_TOKENS_PARAMS.contents,
)
assert response4

View File

@@ -0,0 +1,159 @@
# 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 copy
import pytest
from ... import _transformers as t
from ... import errors
from ... import types
from .. import pytest_helper
from . import constants
_COUNT_TOKENS_PARAMS = types._CountTokensParameters(
model='gemini-2.5-flash',
contents=[t.t_content('Tell me a story in 300 words.')],
)
_COUNT_TOKENS_PARAMS_WITH_SYSTEM_INSTRUCTION = copy.deepcopy(
_COUNT_TOKENS_PARAMS
)
_COUNT_TOKENS_PARAMS_WITH_SYSTEM_INSTRUCTION.config = {
'system_instruction': t.t_content('you are a chatbot.')
}
_COUNT_TOKENS_PARAMS_WITH_TOOLS = copy.deepcopy(_COUNT_TOKENS_PARAMS)
_COUNT_TOKENS_PARAMS_WITH_TOOLS.config = {
'tools': [{'google_search_retrieval': {}}]
}
_COUNT_TOKENS_PARAMS_WITH_GENERATION_CONFIG = copy.deepcopy(
_COUNT_TOKENS_PARAMS
)
_COUNT_TOKENS_PARAMS_WITH_GENERATION_CONFIG.config = {
'generation_config': {'max_output_tokens': 50}
}
_COUNT_TOKENS_PARAMS_VERTEX_CUSTOM_URL = copy.deepcopy(_COUNT_TOKENS_PARAMS)
_COUNT_TOKENS_PARAMS_VERTEX_CUSTOM_URL.config = {
'http_options': constants.VERTEX_HTTP_OPTIONS
}
_COUNT_TOKENS_PARAMS_MLDEV_CUSTOM_URL = copy.deepcopy(_COUNT_TOKENS_PARAMS)
_COUNT_TOKENS_PARAMS_MLDEV_CUSTOM_URL.config = {
'http_options': constants.MLDEV_HTTP_OPTIONS
}
# TODO(b/378952792): MLDev count_tokens needs to merge contents and model
# param into generateContentRequest field.
test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_count_tokens',
parameters=_COUNT_TOKENS_PARAMS,
),
pytest_helper.TestTableItem(
name='test_count_tokens_vertex_custom_url',
parameters=_COUNT_TOKENS_PARAMS_VERTEX_CUSTOM_URL,
exception_if_mldev='404',
),
pytest_helper.TestTableItem(
name='test_count_tokens_mldev_custom_url',
parameters=_COUNT_TOKENS_PARAMS_MLDEV_CUSTOM_URL,
exception_if_vertex='404',
),
pytest_helper.TestTableItem(
name='test_count_tokens_with_system_instruction',
exception_if_mldev='not supported',
parameters=_COUNT_TOKENS_PARAMS_WITH_SYSTEM_INSTRUCTION,
),
pytest_helper.TestTableItem(
name='test_count_tokens_with_tools',
exception_if_mldev='not supported',
parameters=_COUNT_TOKENS_PARAMS_WITH_TOOLS,
),
pytest_helper.TestTableItem(
name='test_count_tokens_with_generation_config',
exception_if_mldev='not supported',
parameters=_COUNT_TOKENS_PARAMS_WITH_GENERATION_CONFIG,
),
]
pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
test_method='models.count_tokens',
test_table=test_table,
)
@pytest.mark.asyncio
async def test_async(client):
response = await client.aio.models.count_tokens(
model=_COUNT_TOKENS_PARAMS.model, contents=_COUNT_TOKENS_PARAMS.contents
)
assert response
def test_different_model_names(client):
if client._api_client.vertexai:
response1 = client.models.count_tokens(
model='gemini-2.5-flash', contents=_COUNT_TOKENS_PARAMS.contents
)
assert response1
response3 = client.models.count_tokens(
model='publishers/google/models/gemini-2.5-flash',
contents=_COUNT_TOKENS_PARAMS.contents,
)
assert response3
response4 = client.models.count_tokens(
model='projects/vertexsdk/locations/us-central1/publishers/google/models/gemini-2.5-flash',
contents=_COUNT_TOKENS_PARAMS.contents,
)
assert response4
else:
response1 = client.models.count_tokens(
model='gemini-2.5-flash', contents=_COUNT_TOKENS_PARAMS.contents
)
assert response1
response2 = client.models.count_tokens(
model='models/gemini-2.5-flash', contents=_COUNT_TOKENS_PARAMS.contents
)
assert response2
def test_extra_body(client):
config = {
'http_options': {
'extra_body': {
'systemInstruction': {
'parts': [{'text': 'you are a chatbot.'}],
'role': 'user',
}
}
}
}
if client._api_client.vertexai:
response = client.models.count_tokens(
model=_COUNT_TOKENS_PARAMS.model,
contents=_COUNT_TOKENS_PARAMS.contents,
config=config,
)
assert response.total_tokens
else:
with pytest.raises(errors.ClientError):
client.models.count_tokens(
model=_COUNT_TOKENS_PARAMS.model,
contents=_COUNT_TOKENS_PARAMS.contents,
config=config,
)

Some files were not shown because too many files have changed in this diff Show More