Files
google--adk-python/tests/unittests/tools/test_google_search_tool.py
T
wehub-resource-sync ec2b666284
Continuous Integration / Pre-commit Linter (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.10) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.11) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.12) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.12) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.14) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Has been cancelled
Copybara PR Handler / close-imported-pr (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 13:25:13 +08:00

539 lines
17 KiB
Python

# Copyright 2026 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 GoogleSearchTool."""
from google.adk.agents.invocation_context import InvocationContext
from google.adk.agents.sequential_agent import SequentialAgent
from google.adk.models.llm_request import LlmRequest
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.adk.tools.google_search_tool import google_search
from google.adk.tools.google_search_tool import GoogleSearchTool
from google.adk.tools.tool_context import ToolContext
from google.genai import types
import pytest
async def _create_tool_context() -> ToolContext:
session_service = InMemorySessionService()
session = await session_service.create_session(
app_name='test_app', user_id='test_user'
)
agent = SequentialAgent(name='test_agent')
invocation_context = InvocationContext(
invocation_id='invocation_id',
agent=agent,
session=session,
session_service=session_service,
)
return ToolContext(invocation_context=invocation_context)
class TestGoogleSearchTool:
"""Test the GoogleSearchTool class."""
def test_init(self):
"""Test initialization of GoogleSearchTool."""
tool = GoogleSearchTool()
assert tool.name == 'google_search'
assert tool.description == 'google_search'
def test_google_search_singleton(self):
"""Test that google_search is a singleton instance."""
assert isinstance(google_search, GoogleSearchTool)
assert google_search.name == 'google_search'
@pytest.mark.asyncio
async def test_process_llm_request_with_gemini_1_model(self):
"""Test processing LLM request with Gemini 1.x model."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
llm_request = LlmRequest(
model='gemini-1.5-flash', config=types.GenerateContentConfig()
)
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
assert llm_request.config.tools is not None
assert len(llm_request.config.tools) == 1
assert llm_request.config.tools[0].google_search_retrieval is not None
@pytest.mark.asyncio
async def test_process_llm_request_with_path_based_gemini_1_model(self):
"""Test processing LLM request with path-based Gemini 1.x model."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
llm_request = LlmRequest(
model='projects/265104255505/locations/us-central1/publishers/google/models/gemini-1.5-flash',
config=types.GenerateContentConfig(),
)
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
assert llm_request.config.tools is not None
assert len(llm_request.config.tools) == 1
assert llm_request.config.tools[0].google_search_retrieval is not None
@pytest.mark.asyncio
async def test_process_llm_request_with_gemini_1_0_model(self):
"""Test processing LLM request with Gemini 1.0 model."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
llm_request = LlmRequest(
model='gemini-1.0-pro', config=types.GenerateContentConfig()
)
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
assert llm_request.config.tools is not None
assert len(llm_request.config.tools) == 1
assert llm_request.config.tools[0].google_search_retrieval is not None
@pytest.mark.asyncio
async def test_process_llm_request_with_gemini_2_model(self):
"""Test processing LLM request with Gemini 2.x model."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
llm_request = LlmRequest(
model='gemini-2.5-flash', config=types.GenerateContentConfig()
)
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
assert llm_request.config.tools is not None
assert len(llm_request.config.tools) == 1
assert llm_request.config.tools[0].google_search is not None
@pytest.mark.asyncio
async def test_process_llm_request_with_path_based_gemini_2_model(self):
"""Test processing LLM request with path-based Gemini 2.x model."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
llm_request = LlmRequest(
model='projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.5-flash',
config=types.GenerateContentConfig(),
)
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
assert llm_request.config.tools is not None
assert len(llm_request.config.tools) == 1
assert llm_request.config.tools[0].google_search is not None
@pytest.mark.asyncio
async def test_process_llm_request_with_gemini_2_5_model(self):
"""Test processing LLM request with Gemini 2.5 model."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
llm_request = LlmRequest(
model='gemini-2.5-pro', config=types.GenerateContentConfig()
)
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
assert llm_request.config.tools is not None
assert len(llm_request.config.tools) == 1
assert llm_request.config.tools[0].google_search is not None
@pytest.mark.asyncio
async def test_process_llm_request_with_gemini_1_model_and_existing_tools_raises_error(
self,
):
"""Test that Gemini 1.x model with existing tools raises ValueError."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
existing_tool = types.Tool(
function_declarations=[
types.FunctionDeclaration(name='test_function', description='test')
]
)
llm_request = LlmRequest(
model='gemini-1.5-flash',
config=types.GenerateContentConfig(tools=[existing_tool]),
)
with pytest.raises(
ValueError,
match=(
'Google search tool cannot be used with other tools in Gemini 1.x'
),
):
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
@pytest.mark.asyncio
async def test_process_llm_request_with_path_based_gemini_1_model_and_existing_tools_raises_error(
self,
):
"""Test that path-based Gemini 1.x model with existing tools raises ValueError."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
existing_tool = types.Tool(
function_declarations=[
types.FunctionDeclaration(name='test_function', description='test')
]
)
llm_request = LlmRequest(
model='projects/265104255505/locations/us-central1/publishers/google/models/gemini-1.5-pro-preview',
config=types.GenerateContentConfig(tools=[existing_tool]),
)
with pytest.raises(
ValueError,
match=(
'Google search tool cannot be used with other tools in Gemini 1.x'
),
):
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
@pytest.mark.asyncio
async def test_process_llm_request_with_gemini_2_model_and_existing_tools_succeeds(
self,
):
"""Test that Gemini 2.x model with existing tools succeeds."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
existing_tool = types.Tool(
function_declarations=[
types.FunctionDeclaration(name='test_function', description='test')
]
)
llm_request = LlmRequest(
model='gemini-2.5-flash',
config=types.GenerateContentConfig(tools=[existing_tool]),
)
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
assert llm_request.config.tools is not None
assert len(llm_request.config.tools) == 2
assert llm_request.config.tools[0] == existing_tool
assert llm_request.config.tools[1].google_search is not None
@pytest.mark.asyncio
async def test_process_llm_request_with_provider_prefixed_gemini_model(
self,
):
"""Test processing LLM request with provider-prefixed Gemini model."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
llm_request = LlmRequest(
model='openrouter/google/gemini-2.5-pro:online',
config=types.GenerateContentConfig(),
)
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
assert llm_request.config.tools is not None
assert len(llm_request.config.tools) == 1
assert llm_request.config.tools[0].google_search is not None
@pytest.mark.asyncio
async def test_process_llm_request_with_non_gemini_model_raises_error(self):
"""Test that non-Gemini model raises ValueError."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
llm_request = LlmRequest(
model='claude-3-sonnet', config=types.GenerateContentConfig()
)
with pytest.raises(
ValueError,
match='Google search tool is not supported for model claude-3-sonnet',
):
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
@pytest.mark.asyncio
async def test_process_llm_request_with_non_gemini_model_and_disabled_check(
self, monkeypatch
):
"""Test non-Gemini model can pass when model-id check is disabled."""
monkeypatch.setenv('ADK_DISABLE_GEMINI_MODEL_ID_CHECK', 'true')
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
llm_request = LlmRequest(
model='internal-model-v1', config=types.GenerateContentConfig()
)
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
assert llm_request.config.tools is not None
assert len(llm_request.config.tools) == 1
assert llm_request.config.tools[0].google_search is not None
@pytest.mark.asyncio
async def test_process_llm_request_with_path_based_non_gemini_model_raises_error(
self,
):
"""Test that path-based non-Gemini model raises ValueError."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
non_gemini_path = 'projects/265104255505/locations/us-central1/publishers/google/models/claude-3-sonnet'
llm_request = LlmRequest(
model=non_gemini_path, config=types.GenerateContentConfig()
)
with pytest.raises(
ValueError,
match=(
f'Google search tool is not supported for model {non_gemini_path}'
),
):
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
@pytest.mark.asyncio
async def test_process_llm_request_with_none_model_raises_error(self):
"""Test that None model raises ValueError."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
llm_request = LlmRequest(model=None, config=types.GenerateContentConfig())
with pytest.raises(
ValueError, match='Google search tool is not supported for model None'
):
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
@pytest.mark.asyncio
async def test_process_llm_request_with_empty_model_raises_error(self):
"""Test that empty model raises ValueError."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
llm_request = LlmRequest(model='', config=types.GenerateContentConfig())
with pytest.raises(
ValueError, match='Google search tool is not supported for model '
):
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
@pytest.mark.asyncio
async def test_process_llm_request_with_no_config(self):
"""Test processing LLM request with None config."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
llm_request = LlmRequest(model='gemini-2.5-flash')
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
assert llm_request.config is not None
assert llm_request.config.tools is not None
assert len(llm_request.config.tools) == 1
assert llm_request.config.tools[0].google_search is not None
@pytest.mark.asyncio
async def test_process_llm_request_with_none_tools(self):
"""Test processing LLM request with None tools."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
llm_request = LlmRequest(
model='gemini-2.5-flash', config=types.GenerateContentConfig(tools=None)
)
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
assert llm_request.config.tools is not None
assert len(llm_request.config.tools) == 1
assert llm_request.config.tools[0].google_search is not None
@pytest.mark.asyncio
async def test_process_llm_request_edge_cases(self):
"""Test edge cases for model name validation."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
# Test with model names that contain gemini but don't start with it
edge_cases = [
'my-gemini-1.5-model',
'custom-gemini-2.5-flash',
'projects/265104255505/locations/us-central1/publishers/gemini/models/claude-3-sonnet',
]
for model in edge_cases:
llm_request = LlmRequest(
model=model, config=types.GenerateContentConfig()
)
with pytest.raises(
ValueError,
match=f'Google search tool is not supported for model {model}',
):
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
@pytest.mark.asyncio
async def test_process_llm_request_gemini_version_specifics(self):
"""Test specific Gemini version behaviors."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
# Test various Gemini versions
gemini_1_models = [
'gemini-1.0-pro',
'gemini-1.5-flash',
'gemini-1.5-pro',
'gemini-1.9-experimental',
]
gemini_2_models = [
'gemini-2.0-pro',
'gemini-2.5-flash',
'gemini-2.5-pro',
]
# Test Gemini 1.x models use google_search_retrieval
for model in gemini_1_models:
llm_request = LlmRequest(
model=model, config=types.GenerateContentConfig()
)
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
assert llm_request.config.tools is not None
assert len(llm_request.config.tools) == 1
assert llm_request.config.tools[0].google_search_retrieval is not None
assert llm_request.config.tools[0].google_search is None
# Test Gemini 2.x models use google_search
for model in gemini_2_models:
llm_request = LlmRequest(
model=model, config=types.GenerateContentConfig()
)
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
assert llm_request.config.tools is not None
assert len(llm_request.config.tools) == 1
assert llm_request.config.tools[0].google_search is not None
assert llm_request.config.tools[0].google_search_retrieval is None
@pytest.mark.asyncio
@pytest.mark.parametrize(
(
'tool_model',
'request_model',
'expected_model',
),
[
(
'gemini-2.5-flash-lite',
'gemini-2.5-flash',
'gemini-2.5-flash-lite',
),
(
None,
'gemini-2.5-flash',
'gemini-2.5-flash',
),
],
ids=['with_custom_model', 'without_custom_model'],
)
async def test_process_llm_request_custom_model_behavior(
self,
tool_model,
request_model,
expected_model,
):
"""Tests custom model parameter behavior in process_llm_request."""
tool = GoogleSearchTool(model=tool_model)
tool_context = await _create_tool_context()
llm_request = LlmRequest(
model=request_model, config=types.GenerateContentConfig()
)
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
assert llm_request.model == expected_model
assert llm_request.config.tools is not None
assert len(llm_request.config.tools) == 1
@pytest.mark.asyncio
async def test_process_llm_request_managed_agent_no_model(self):
"""Managed-agent requests resolve google_search even with no model."""
tool = GoogleSearchTool()
tool_context = await _create_tool_context()
llm_request = LlmRequest(
model=None,
config=types.GenerateContentConfig(),
)
llm_request._is_managed_agent = True
await tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
assert llm_request.config.tools is not None
assert len(llm_request.config.tools) == 1
assert llm_request.config.tools[0].google_search is not None