encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

Moves unittests to root folder and adds github action to run unit tests. (#72)

Browse Source 
* Move unit tests to root package.

* Adds deps to "test" extra, and mark two broken tests in tests/unittests/auth/test_auth_handler.py

* Adds github workflow

* minor fix in lite_llm.py for python 3.9.

* format pyproject.toml
This commit is contained in:
Jack Sun
2025-04-11 08:25:59 -07:00
committed by GitHub
parent 59117b9b96
commit 05142a07cc
66 changed files with 50 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/models/lite_llm.py
+3 -1
View File
@@ -183,7 +183,9 @@ def _content_to_message_param(
)
def _get_content(parts: Iterable[types.Part]) -> OpenAIMessageContent | str:
def _get_content(
parts: Iterable[types.Part],
) -> Union[OpenAIMessageContent, str]:
"""Converts a list of parts to litellm content.
Args:
src/google/adk/tests/unittests/__init__.py
-14
View File
@@ -1,14 +0,0 @@
# 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.
src/google/adk/tests/unittests/agents/__init__.py
-14
View File
@@ -1,14 +0,0 @@
# 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.
src/google/adk/tests/unittests/agents/test_base_agent.py
-407
View File
@@ -1,407 +0,0 @@
# 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.
"""Testings for the BaseAgent."""
from typing import AsyncGenerator
from typing import Optional
from typing import Union
from google.adk.agents.base_agent import BaseAgent
from google.adk.agents.callback_context import CallbackContext
from google.adk.agents.invocation_context import InvocationContext
from google.adk.events import Event
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.genai import types
import pytest
import pytest_mock
from typing_extensions import override
def _before_agent_callback_noop(callback_context: CallbackContext) -> None:
pass
def _before_agent_callback_bypass_agent(
callback_context: CallbackContext,
) -> types.Content:
return types.Content(parts=[types.Part(text='agent run is bypassed.')])
def _after_agent_callback_noop(callback_context: CallbackContext) -> None:
pass
def _after_agent_callback_append_agent_reply(
callback_context: CallbackContext,
) -> types.Content:
return types.Content(
parts=[types.Part(text='Agent reply from after agent callback.')]
)
class _IncompleteAgent(BaseAgent):
pass
class _TestingAgent(BaseAgent):
@override
async def _run_async_impl(
self, ctx: InvocationContext
) -> AsyncGenerator[Event, None]:
yield Event(
author=self.name,
branch=ctx.branch,
invocation_id=ctx.invocation_id,
content=types.Content(parts=[types.Part(text='Hello, world!')]),
)
@override
async def _run_live_impl(
self, ctx: InvocationContext
) -> AsyncGenerator[Event, None]:
yield Event(
author=self.name,
invocation_id=ctx.invocation_id,
branch=ctx.branch,
content=types.Content(parts=[types.Part(text='Hello, live!')]),
)
def _create_parent_invocation_context(
test_name: str, agent: BaseAgent, branch: Optional[str] = None
) -> InvocationContext:
session_service = InMemorySessionService()
session = session_service.create_session(
app_name='test_app', user_id='test_user'
)
return InvocationContext(
invocation_id=f'{test_name}_invocation_id',
branch=branch,
agent=agent,
session=session,
session_service=session_service,
)
def test_invalid_agent_name():
with pytest.raises(ValueError):
_ = _TestingAgent(name='not an identifier')
@pytest.mark.asyncio
async def test_run_async(request: pytest.FixtureRequest):
agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
parent_ctx = _create_parent_invocation_context(
request.function.__name__, agent
)
events = [e async for e in agent.run_async(parent_ctx)]
assert len(events) == 1
assert events[0].author == agent.name
assert events[0].content.parts[0].text == 'Hello, world!'
@pytest.mark.asyncio
async def test_run_async_with_branch(request: pytest.FixtureRequest):
agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
parent_ctx = _create_parent_invocation_context(
request.function.__name__, agent, branch='parent_branch'
)
events = [e async for e in agent.run_async(parent_ctx)]
assert len(events) == 1
assert events[0].author == agent.name
assert events[0].content.parts[0].text == 'Hello, world!'
assert events[0].branch.endswith(agent.name)
@pytest.mark.asyncio
async def test_run_async_before_agent_callback_noop(
request: pytest.FixtureRequest,
mocker: pytest_mock.MockerFixture,
) -> Union[types.Content, None]:
# Arrange
agent = _TestingAgent(
name=f'{request.function.__name__}_test_agent',
before_agent_callback=_before_agent_callback_noop,
)
parent_ctx = _create_parent_invocation_context(
request.function.__name__, agent
)
spy_run_async_impl = mocker.spy(agent, BaseAgent._run_async_impl.__name__)
spy_before_agent_callback = mocker.spy(agent, 'before_agent_callback')
# Act
_ = [e async for e in agent.run_async(parent_ctx)]
# Assert
spy_before_agent_callback.assert_called_once()
_, kwargs = spy_before_agent_callback.call_args
assert 'callback_context' in kwargs
assert isinstance(kwargs['callback_context'], CallbackContext)
spy_run_async_impl.assert_called_once()
@pytest.mark.asyncio
async def test_run_async_before_agent_callback_bypass_agent(
request: pytest.FixtureRequest,
mocker: pytest_mock.MockerFixture,
):
# Arrange
agent = _TestingAgent(
name=f'{request.function.__name__}_test_agent',
before_agent_callback=_before_agent_callback_bypass_agent,
)
parent_ctx = _create_parent_invocation_context(
request.function.__name__, agent
)
spy_run_async_impl = mocker.spy(agent, BaseAgent._run_async_impl.__name__)
spy_before_agent_callback = mocker.spy(agent, 'before_agent_callback')
# Act
events = [e async for e in agent.run_async(parent_ctx)]
# Assert
spy_before_agent_callback.assert_called_once()
spy_run_async_impl.assert_not_called()
assert len(events) == 1
assert events[0].content.parts[0].text == 'agent run is bypassed.'
@pytest.mark.asyncio
async def test_run_async_after_agent_callback_noop(
request: pytest.FixtureRequest,
mocker: pytest_mock.MockerFixture,
):
# Arrange
agent = _TestingAgent(
name=f'{request.function.__name__}_test_agent',
after_agent_callback=_after_agent_callback_noop,
)
parent_ctx = _create_parent_invocation_context(
request.function.__name__, agent
)
spy_after_agent_callback = mocker.spy(agent, 'after_agent_callback')
# Act
events = [e async for e in agent.run_async(parent_ctx)]
# Assert
spy_after_agent_callback.assert_called_once()
_, kwargs = spy_after_agent_callback.call_args
assert 'callback_context' in kwargs
assert isinstance(kwargs['callback_context'], CallbackContext)
assert len(events) == 1
@pytest.mark.asyncio
async def test_run_async_after_agent_callback_append_reply(
request: pytest.FixtureRequest,
):
# Arrange
agent = _TestingAgent(
name=f'{request.function.__name__}_test_agent',
after_agent_callback=_after_agent_callback_append_agent_reply,
)
parent_ctx = _create_parent_invocation_context(
request.function.__name__, agent
)
# Act
events = [e async for e in agent.run_async(parent_ctx)]
# Assert
assert len(events) == 2
assert events[1].author == agent.name
assert (
events[1].content.parts[0].text
== 'Agent reply from after agent callback.'
)
@pytest.mark.asyncio
async def test_run_async_incomplete_agent(request: pytest.FixtureRequest):
agent = _IncompleteAgent(name=f'{request.function.__name__}_test_agent')
parent_ctx = _create_parent_invocation_context(
request.function.__name__, agent
)
with pytest.raises(NotImplementedError):
[e async for e in agent.run_async(parent_ctx)]
@pytest.mark.asyncio
async def test_run_live(request: pytest.FixtureRequest):
agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
parent_ctx = _create_parent_invocation_context(
request.function.__name__, agent
)
events = [e async for e in agent.run_live(parent_ctx)]
assert len(events) == 1
assert events[0].author == agent.name
assert events[0].content.parts[0].text == 'Hello, live!'
@pytest.mark.asyncio
async def test_run_live_with_branch(request: pytest.FixtureRequest):
agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
parent_ctx = _create_parent_invocation_context(
request.function.__name__, agent, branch='parent_branch'
)
events = [e async for e in agent.run_live(parent_ctx)]
assert len(events) == 1
assert events[0].author == agent.name
assert events[0].content.parts[0].text == 'Hello, live!'
assert events[0].branch.endswith(agent.name)
@pytest.mark.asyncio
async def test_run_live_incomplete_agent(request: pytest.FixtureRequest):
agent = _IncompleteAgent(name=f'{request.function.__name__}_test_agent')
parent_ctx = _create_parent_invocation_context(
request.function.__name__, agent
)
with pytest.raises(NotImplementedError):
[e async for e in agent.run_live(parent_ctx)]
def test_set_parent_agent_for_sub_agents(request: pytest.FixtureRequest):
sub_agents: list[BaseAgent] = [
_TestingAgent(name=f'{request.function.__name__}_sub_agent_1'),
_TestingAgent(name=f'{request.function.__name__}_sub_agent_2'),
]
parent = _TestingAgent(
name=f'{request.function.__name__}_parent',
sub_agents=sub_agents,
)
for sub_agent in sub_agents:
assert sub_agent.parent_agent == parent
def test_find_agent(request: pytest.FixtureRequest):
grand_sub_agent_1 = _TestingAgent(
name=f'{request.function.__name__}__grand_sub_agent_1'
)
grand_sub_agent_2 = _TestingAgent(
name=f'{request.function.__name__}__grand_sub_agent_2'
)
sub_agent_1 = _TestingAgent(
name=f'{request.function.__name__}_sub_agent_1',
sub_agents=[grand_sub_agent_1],
)
sub_agent_2 = _TestingAgent(
name=f'{request.function.__name__}_sub_agent_2',
sub_agents=[grand_sub_agent_2],
)
parent = _TestingAgent(
name=f'{request.function.__name__}_parent',
sub_agents=[sub_agent_1, sub_agent_2],
)
assert parent.find_agent(parent.name) == parent
assert parent.find_agent(sub_agent_1.name) == sub_agent_1
assert parent.find_agent(sub_agent_2.name) == sub_agent_2
assert parent.find_agent(grand_sub_agent_1.name) == grand_sub_agent_1
assert parent.find_agent(grand_sub_agent_2.name) == grand_sub_agent_2
assert sub_agent_1.find_agent(grand_sub_agent_1.name) == grand_sub_agent_1
assert sub_agent_1.find_agent(grand_sub_agent_2.name) is None
assert sub_agent_2.find_agent(grand_sub_agent_1.name) is None
assert sub_agent_2.find_agent(sub_agent_2.name) == sub_agent_2
assert parent.find_agent('not_exist') is None
def test_find_sub_agent(request: pytest.FixtureRequest):
grand_sub_agent_1 = _TestingAgent(
name=f'{request.function.__name__}__grand_sub_agent_1'
)
grand_sub_agent_2 = _TestingAgent(
name=f'{request.function.__name__}__grand_sub_agent_2'
)
sub_agent_1 = _TestingAgent(
name=f'{request.function.__name__}_sub_agent_1',
sub_agents=[grand_sub_agent_1],
)
sub_agent_2 = _TestingAgent(
name=f'{request.function.__name__}_sub_agent_2',
sub_agents=[grand_sub_agent_2],
)
parent = _TestingAgent(
name=f'{request.function.__name__}_parent',
sub_agents=[sub_agent_1, sub_agent_2],
)
assert parent.find_sub_agent(sub_agent_1.name) == sub_agent_1
assert parent.find_sub_agent(sub_agent_2.name) == sub_agent_2
assert parent.find_sub_agent(grand_sub_agent_1.name) == grand_sub_agent_1
assert parent.find_sub_agent(grand_sub_agent_2.name) == grand_sub_agent_2
assert sub_agent_1.find_sub_agent(grand_sub_agent_1.name) == grand_sub_agent_1
assert sub_agent_1.find_sub_agent(grand_sub_agent_2.name) is None
assert sub_agent_2.find_sub_agent(grand_sub_agent_1.name) is None
assert sub_agent_2.find_sub_agent(grand_sub_agent_2.name) == grand_sub_agent_2
assert parent.find_sub_agent(parent.name) is None
assert parent.find_sub_agent('not_exist') is None
def test_root_agent(request: pytest.FixtureRequest):
grand_sub_agent_1 = _TestingAgent(
name=f'{request.function.__name__}__grand_sub_agent_1'
)
grand_sub_agent_2 = _TestingAgent(
name=f'{request.function.__name__}__grand_sub_agent_2'
)
sub_agent_1 = _TestingAgent(
name=f'{request.function.__name__}_sub_agent_1',
sub_agents=[grand_sub_agent_1],
)
sub_agent_2 = _TestingAgent(
name=f'{request.function.__name__}_sub_agent_2',
sub_agents=[grand_sub_agent_2],
)
parent = _TestingAgent(
name=f'{request.function.__name__}_parent',
sub_agents=[sub_agent_1, sub_agent_2],
)
assert parent.root_agent == parent
assert sub_agent_1.root_agent == parent
assert sub_agent_2.root_agent == parent
assert grand_sub_agent_1.root_agent == parent
assert grand_sub_agent_2.root_agent == parent
def test_set_parent_agent_for_sub_agent_twice(
request: pytest.FixtureRequest,
):
sub_agent = _TestingAgent(name=f'{request.function.__name__}_sub_agent')
_ = _TestingAgent(
name=f'{request.function.__name__}_parent_1',
sub_agents=[sub_agent],
)
with pytest.raises(ValueError):
_ = _TestingAgent(
name=f'{request.function.__name__}_parent_2',
sub_agents=[sub_agent],
)
src/google/adk/tests/unittests/agents/test_langgraph_agent.py
-191
View File
@@ -1,191 +0,0 @@
# 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.mock import MagicMock
from google.adk.agents.invocation_context import InvocationContext
from google.adk.agents.langgraph_agent import LangGraphAgent
from google.adk.events import Event
from google.genai import types
from langchain_core.messages import AIMessage
from langchain_core.messages import HumanMessage
from langchain_core.messages import SystemMessage
from langgraph.graph.graph import CompiledGraph
import pytest
@pytest.mark.parametrize(
"checkpointer_value, events_list, expected_messages",
[
(
MagicMock(),
[
Event(
invocation_id="test_invocation_id",
author="user",
content=types.Content(
role="user",
parts=[types.Part.from_text(text="test prompt")],
),
),
Event(
invocation_id="test_invocation_id",
author="root_agent",
content=types.Content(
role="model",
parts=[types.Part.from_text(text="(some delegation)")],
),
),
],
[
SystemMessage(content="test system prompt"),
HumanMessage(content="test prompt"),
],
),
(
None,
[
Event(
invocation_id="test_invocation_id",
author="user",
content=types.Content(
role="user",
parts=[types.Part.from_text(text="user prompt 1")],
),
),
Event(
invocation_id="test_invocation_id",
author="root_agent",
content=types.Content(
role="model",
parts=[
types.Part.from_text(text="root agent response")
],
),
),
Event(
invocation_id="test_invocation_id",
author="weather_agent",
content=types.Content(
role="model",
parts=[
types.Part.from_text(text="weather agent response")
],
),
),
Event(
invocation_id="test_invocation_id",
author="user",
content=types.Content(
role="user",
parts=[types.Part.from_text(text="user prompt 2")],
),
),
],
[
SystemMessage(content="test system prompt"),
HumanMessage(content="user prompt 1"),
AIMessage(content="weather agent response"),
HumanMessage(content="user prompt 2"),
],
),
(
MagicMock(),
[
Event(
invocation_id="test_invocation_id",
author="user",
content=types.Content(
role="user",
parts=[types.Part.from_text(text="user prompt 1")],
),
),
Event(
invocation_id="test_invocation_id",
author="root_agent",
content=types.Content(
role="model",
parts=[
types.Part.from_text(text="root agent response")
],
),
),
Event(
invocation_id="test_invocation_id",
author="weather_agent",
content=types.Content(
role="model",
parts=[
types.Part.from_text(text="weather agent response")
],
),
),
Event(
invocation_id="test_invocation_id",
author="user",
content=types.Content(
role="user",
parts=[types.Part.from_text(text="user prompt 2")],
),
),
],
[
SystemMessage(content="test system prompt"),
HumanMessage(content="user prompt 2"),
],
),
],
)
@pytest.mark.asyncio
async def test_langgraph_agent(
checkpointer_value, events_list, expected_messages
):
mock_graph = MagicMock(spec=CompiledGraph)
mock_graph_state = MagicMock()
mock_graph_state.values = {}
mock_graph.get_state.return_value = mock_graph_state
mock_graph.checkpointer = checkpointer_value
mock_graph.invoke.return_value = {
"messages": [AIMessage(content="test response")]
}
mock_parent_context = MagicMock(spec=InvocationContext)
mock_session = MagicMock()
mock_parent_context.session = mock_session
mock_parent_context.branch = "parent_agent"
mock_parent_context.end_invocation = False
mock_session.events = events_list
mock_parent_context.invocation_id = "test_invocation_id"
mock_parent_context.model_copy.return_value = mock_parent_context
weather_agent = LangGraphAgent(
name="weather_agent",
description="A agent that answers weather questions",
instruction="test system prompt",
graph=mock_graph,
)
result_event = None
async for event in weather_agent.run_async(mock_parent_context):
result_event = event
assert result_event.author == "weather_agent"
assert result_event.content.parts[0].text == "test response"
mock_graph.invoke.assert_called_once()
mock_graph.invoke.assert_called_with(
{"messages": expected_messages},
{"configurable": {"thread_id": mock_session.id}},
)
src/google/adk/tests/unittests/agents/test_llm_agent_callbacks.py
-138
View File
@@ -1,138 +0,0 @@
# 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 typing import Any
from typing import Optional
from google.adk.agents.callback_context import CallbackContext
from google.adk.agents.llm_agent import Agent
from google.adk.models import LlmRequest
from google.adk.models import LlmResponse
from google.genai import types
from pydantic import BaseModel
import pytest
from .. import utils
class MockBeforeModelCallback(BaseModel):
mock_response: str
def __call__(
self,
callback_context: CallbackContext,
llm_request: LlmRequest,
) -> LlmResponse:
return LlmResponse(
content=utils.ModelContent(
[types.Part.from_text(text=self.mock_response)]
)
)
class MockAfterModelCallback(BaseModel):
mock_response: str
def __call__(
self,
callback_context: CallbackContext,
llm_response: LlmResponse,
) -> LlmResponse:
return LlmResponse(
content=utils.ModelContent(
[types.Part.from_text(text=self.mock_response)]
)
)
def noop_callback(**kwargs) -> Optional[LlmResponse]:
pass
@pytest.mark.asyncio
async def test_before_model_callback():
responses = ['model_response']
mock_model = utils.MockModel.create(responses=responses)
agent = Agent(
name='root_agent',
model=mock_model,
before_model_callback=MockBeforeModelCallback(
mock_response='before_model_callback'
),
)
runner = utils.TestInMemoryRunner(agent)
assert utils.simplify_events(
await runner.run_async_with_new_session('test')
) == [
('root_agent', 'before_model_callback'),
]
@pytest.mark.asyncio
async def test_before_model_callback_noop():
responses = ['model_response']
mock_model = utils.MockModel.create(responses=responses)
agent = Agent(
name='root_agent',
model=mock_model,
before_model_callback=noop_callback,
)
runner = utils.TestInMemoryRunner(agent)
assert utils.simplify_events(
await runner.run_async_with_new_session('test')
) == [
('root_agent', 'model_response'),
]
@pytest.mark.asyncio
async def test_before_model_callback_end():
responses = ['model_response']
mock_model = utils.MockModel.create(responses=responses)
agent = Agent(
name='root_agent',
model=mock_model,
before_model_callback=MockBeforeModelCallback(
mock_response='before_model_callback',
),
)
runner = utils.TestInMemoryRunner(agent)
assert utils.simplify_events(
await runner.run_async_with_new_session('test')
) == [
('root_agent', 'before_model_callback'),
]
@pytest.mark.asyncio
async def test_after_model_callback():
responses = ['model_response']
mock_model = utils.MockModel.create(responses=responses)
agent = Agent(
name='root_agent',
model=mock_model,
after_model_callback=MockAfterModelCallback(
mock_response='after_model_callback'
),
)
runner = utils.TestInMemoryRunner(agent)
assert utils.simplify_events(
await runner.run_async_with_new_session('test')
) == [
('root_agent', 'after_model_callback'),
]
src/google/adk/tests/unittests/agents/test_llm_agent_fields.py
-231
View File
@@ -1,231 +0,0 @@
# 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 canonical_xxx fields in LlmAgent."""
from typing import Any
from typing import Optional
from google.adk.agents.callback_context import CallbackContext
from google.adk.agents.invocation_context import InvocationContext
from google.adk.agents.llm_agent import LlmAgent
from google.adk.agents.loop_agent import LoopAgent
from google.adk.agents.readonly_context import ReadonlyContext
from google.adk.models.llm_request import LlmRequest
from google.adk.models.registry import LLMRegistry
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.genai import types
from pydantic import BaseModel
import pytest
def _create_readonly_context(
agent: LlmAgent, state: Optional[dict[str, Any]] = None
) -> ReadonlyContext:
session_service = InMemorySessionService()
session = session_service.create_session(
app_name='test_app', user_id='test_user', state=state
)
invocation_context = InvocationContext(
invocation_id='test_id',
agent=agent,
session=session,
session_service=session_service,
)
return ReadonlyContext(invocation_context)
def test_canonical_model_empty():
agent = LlmAgent(name='test_agent')
with pytest.raises(ValueError):
_ = agent.canonical_model
def test_canonical_model_str():
agent = LlmAgent(name='test_agent', model='gemini-pro')
assert agent.canonical_model.model == 'gemini-pro'
def test_canonical_model_llm():
llm = LLMRegistry.new_llm('gemini-pro')
agent = LlmAgent(name='test_agent', model=llm)
assert agent.canonical_model == llm
def test_canonical_model_inherit():
sub_agent = LlmAgent(name='sub_agent')
parent_agent = LlmAgent(
name='parent_agent', model='gemini-pro', sub_agents=[sub_agent]
)
assert sub_agent.canonical_model == parent_agent.canonical_model
def test_canonical_instruction_str():
agent = LlmAgent(name='test_agent', instruction='instruction')
ctx = _create_readonly_context(agent)
assert agent.canonical_instruction(ctx) == 'instruction'
def test_canonical_instruction():
def _instruction_provider(ctx: ReadonlyContext) -> str:
return f'instruction: {ctx.state["state_var"]}'
agent = LlmAgent(name='test_agent', instruction=_instruction_provider)
ctx = _create_readonly_context(agent, state={'state_var': 'state_value'})
assert agent.canonical_instruction(ctx) == 'instruction: state_value'
def test_canonical_global_instruction_str():
agent = LlmAgent(name='test_agent', global_instruction='global instruction')
ctx = _create_readonly_context(agent)
assert agent.canonical_global_instruction(ctx) == 'global instruction'
def test_canonical_global_instruction():
def _global_instruction_provider(ctx: ReadonlyContext) -> str:
return f'global instruction: {ctx.state["state_var"]}'
agent = LlmAgent(
name='test_agent', global_instruction=_global_instruction_provider
)
ctx = _create_readonly_context(agent, state={'state_var': 'state_value'})
assert (
agent.canonical_global_instruction(ctx)
== 'global instruction: state_value'
)
def test_output_schema_will_disable_transfer(caplog: pytest.LogCaptureFixture):
with caplog.at_level('WARNING'):
class Schema(BaseModel):
pass
agent = LlmAgent(
name='test_agent',
output_schema=Schema,
)
# Transfer is automatically disabled
assert agent.disallow_transfer_to_parent
assert agent.disallow_transfer_to_peers
assert (
'output_schema cannot co-exist with agent transfer configurations.'
in caplog.text
)
def test_output_schema_with_sub_agents_will_throw():
class Schema(BaseModel):
pass
sub_agent = LlmAgent(
name='sub_agent',
)
with pytest.raises(ValueError):
_ = LlmAgent(
name='test_agent',
output_schema=Schema,
sub_agents=[sub_agent],
)
def test_output_schema_with_tools_will_throw():
class Schema(BaseModel):
pass
def _a_tool():
pass
with pytest.raises(ValueError):
_ = LlmAgent(
name='test_agent',
output_schema=Schema,
tools=[_a_tool],
)
def test_before_model_callback():
def _before_model_callback(
callback_context: CallbackContext,
llm_request: LlmRequest,
) -> None:
return None
agent = LlmAgent(
name='test_agent', before_model_callback=_before_model_callback
)
# TODO: add more logic assertions later.
assert agent.before_model_callback is not None
def test_validate_generate_content_config_thinking_config_throw():
with pytest.raises(ValueError):
_ = LlmAgent(
name='test_agent',
generate_content_config=types.GenerateContentConfig(
thinking_config=types.ThinkingConfig()
),
)
def test_validate_generate_content_config_tools_throw():
with pytest.raises(ValueError):
_ = LlmAgent(
name='test_agent',
generate_content_config=types.GenerateContentConfig(
tools=[types.Tool(function_declarations=[])]
),
)
def test_validate_generate_content_config_system_instruction_throw():
with pytest.raises(ValueError):
_ = LlmAgent(
name='test_agent',
generate_content_config=types.GenerateContentConfig(
system_instruction='system instruction'
),
)
def test_validate_generate_content_config_response_schema_throw():
class Schema(BaseModel):
pass
with pytest.raises(ValueError):
_ = LlmAgent(
name='test_agent',
generate_content_config=types.GenerateContentConfig(
response_schema=Schema
),
)
def test_allow_transfer_by_default():
sub_agent = LlmAgent(name='sub_agent')
agent = LlmAgent(name='test_agent', sub_agents=[sub_agent])
assert not agent.disallow_transfer_to_parent
assert not agent.disallow_transfer_to_peers
src/google/adk/tests/unittests/agents/test_loop_agent.py
-136
View File
@@ -1,136 +0,0 @@
# 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.
"""Testings for the SequentialAgent."""
from typing import AsyncGenerator
from google.adk.agents.base_agent import BaseAgent
from google.adk.agents.invocation_context import InvocationContext
from google.adk.agents.loop_agent import LoopAgent
from google.adk.events import Event
from google.adk.events import EventActions
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.genai import types
import pytest
from typing_extensions import override
class _TestingAgent(BaseAgent):
@override
async def _run_async_impl(
self, ctx: InvocationContext
) -> AsyncGenerator[Event, None]:
yield Event(
author=self.name,
invocation_id=ctx.invocation_id,
content=types.Content(
parts=[types.Part(text=f'Hello, async {self.name}!')]
),
)
@override
async def _run_live_impl(
self, ctx: InvocationContext
) -> AsyncGenerator[Event, None]:
yield Event(
author=self.name,
invocation_id=ctx.invocation_id,
content=types.Content(
parts=[types.Part(text=f'Hello, live {self.name}!')]
),
)
class _TestingAgentWithEscalateAction(BaseAgent):
@override
async def _run_async_impl(
self, ctx: InvocationContext
) -> AsyncGenerator[Event, None]:
yield Event(
author=self.name,
invocation_id=ctx.invocation_id,
content=types.Content(
parts=[types.Part(text=f'Hello, async {self.name}!')]
),
actions=EventActions(escalate=True),
)
def _create_parent_invocation_context(
test_name: str, agent: BaseAgent
) -> InvocationContext:
session_service = InMemorySessionService()
session = session_service.create_session(
app_name='test_app', user_id='test_user'
)
return InvocationContext(
invocation_id=f'{test_name}_invocation_id',
agent=agent,
session=session,
session_service=session_service,
)
@pytest.mark.asyncio
async def test_run_async(request: pytest.FixtureRequest):
agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
loop_agent = LoopAgent(
name=f'{request.function.__name__}_test_loop_agent',
max_iterations=2,
sub_agents=[
agent,
],
)
parent_ctx = _create_parent_invocation_context(
request.function.__name__, loop_agent
)
events = [e async for e in loop_agent.run_async(parent_ctx)]
assert len(events) == 2
assert events[0].author == agent.name
assert events[1].author == agent.name
assert events[0].content.parts[0].text == f'Hello, async {agent.name}!'
assert events[1].content.parts[0].text == f'Hello, async {agent.name}!'
@pytest.mark.asyncio
async def test_run_async_with_escalate_action(request: pytest.FixtureRequest):
non_escalating_agent = _TestingAgent(
name=f'{request.function.__name__}_test_non_escalating_agent'
)
escalating_agent = _TestingAgentWithEscalateAction(
name=f'{request.function.__name__}_test_escalating_agent'
)
loop_agent = LoopAgent(
name=f'{request.function.__name__}_test_loop_agent',
sub_agents=[non_escalating_agent, escalating_agent],
)
parent_ctx = _create_parent_invocation_context(
request.function.__name__, loop_agent
)
events = [e async for e in loop_agent.run_async(parent_ctx)]
# Only two events are generated because the sub escalating_agent escalates.
assert len(events) == 2
assert events[0].author == non_escalating_agent.name
assert events[1].author == escalating_agent.name
assert events[0].content.parts[0].text == (
f'Hello, async {non_escalating_agent.name}!'
)
assert events[1].content.parts[0].text == (
f'Hello, async {escalating_agent.name}!'
)
src/google/adk/tests/unittests/agents/test_parallel_agent.py
-92
View File
@@ -1,92 +0,0 @@
# 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 ParallelAgent."""
import asyncio
from typing import AsyncGenerator
from google.adk.agents.base_agent import BaseAgent
from google.adk.agents.invocation_context import InvocationContext
from google.adk.agents.parallel_agent import ParallelAgent
from google.adk.events import Event
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.genai import types
import pytest
from typing_extensions import override
class _TestingAgent(BaseAgent):
delay: float = 0
"""The delay before the agent generates an event."""
@override
async def _run_async_impl(
self, ctx: InvocationContext
) -> AsyncGenerator[Event, None]:
await asyncio.sleep(self.delay)
yield Event(
author=self.name,
branch=ctx.branch,
invocation_id=ctx.invocation_id,
content=types.Content(
parts=[types.Part(text=f'Hello, async {self.name}!')]
),
)
def _create_parent_invocation_context(
test_name: str, agent: BaseAgent
) -> InvocationContext:
session_service = InMemorySessionService()
session = session_service.create_session(
app_name='test_app', user_id='test_user'
)
return InvocationContext(
invocation_id=f'{test_name}_invocation_id',
agent=agent,
session=session,
session_service=session_service,
)
@pytest.mark.asyncio
async def test_run_async(request: pytest.FixtureRequest):
agent1 = _TestingAgent(
name=f'{request.function.__name__}_test_agent_1',
delay=0.5,
)
agent2 = _TestingAgent(name=f'{request.function.__name__}_test_agent_2')
parallel_agent = ParallelAgent(
name=f'{request.function.__name__}_test_parallel_agent',
sub_agents=[
agent1,
agent2,
],
)
parent_ctx = _create_parent_invocation_context(
request.function.__name__, parallel_agent
)
events = [e async for e in parallel_agent.run_async(parent_ctx)]
assert len(events) == 2
# agent2 generates an event first, then agent1. Because they run in parallel
# and agent1 has a delay.
assert events[0].author == agent2.name
assert events[1].author == agent1.name
assert events[0].branch.endswith(agent2.name)
assert events[1].branch.endswith(agent1.name)
assert events[0].content.parts[0].text == f'Hello, async {agent2.name}!'
assert events[1].content.parts[0].text == f'Hello, async {agent1.name}!'
src/google/adk/tests/unittests/agents/test_sequential_agent.py
-114
View File
@@ -1,114 +0,0 @@
# 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.
"""Testings for the SequentialAgent."""
from typing import AsyncGenerator
from google.adk.agents.base_agent import BaseAgent
from google.adk.agents.invocation_context import InvocationContext
from google.adk.agents.sequential_agent import SequentialAgent
from google.adk.events import Event
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.genai import types
import pytest
from typing_extensions import override
class _TestingAgent(BaseAgent):
@override
async def _run_async_impl(
self, ctx: InvocationContext
) -> AsyncGenerator[Event, None]:
yield Event(
author=self.name,
invocation_id=ctx.invocation_id,
content=types.Content(
parts=[types.Part(text=f'Hello, async {self.name}!')]
),
)
@override
async def _run_live_impl(
self, ctx: InvocationContext
) -> AsyncGenerator[Event, None]:
yield Event(
author=self.name,
invocation_id=ctx.invocation_id,
content=types.Content(
parts=[types.Part(text=f'Hello, live {self.name}!')]
),
)
def _create_parent_invocation_context(
test_name: str, agent: BaseAgent
) -> InvocationContext:
session_service = InMemorySessionService()
session = session_service.create_session(
app_name='test_app', user_id='test_user'
)
return InvocationContext(
invocation_id=f'{test_name}_invocation_id',
agent=agent,
session=session,
session_service=session_service,
)
@pytest.mark.asyncio
async def test_run_async(request: pytest.FixtureRequest):
agent_1 = _TestingAgent(name=f'{request.function.__name__}_test_agent_1')
agent_2 = _TestingAgent(name=f'{request.function.__name__}_test_agent_2')
sequential_agent = SequentialAgent(
name=f'{request.function.__name__}_test_agent',
sub_agents=[
agent_1,
agent_2,
],
)
parent_ctx = _create_parent_invocation_context(
request.function.__name__, sequential_agent
)
events = [e async for e in sequential_agent.run_async(parent_ctx)]
assert len(events) == 2
assert events[0].author == agent_1.name
assert events[1].author == agent_2.name
assert events[0].content.parts[0].text == f'Hello, async {agent_1.name}!'
assert events[1].content.parts[0].text == f'Hello, async {agent_2.name}!'
@pytest.mark.asyncio
async def test_run_live(request: pytest.FixtureRequest):
agent_1 = _TestingAgent(name=f'{request.function.__name__}_test_agent_1')
agent_2 = _TestingAgent(name=f'{request.function.__name__}_test_agent_2')
sequential_agent = SequentialAgent(
name=f'{request.function.__name__}_test_agent',
sub_agents=[
agent_1,
agent_2,
],
)
parent_ctx = _create_parent_invocation_context(
request.function.__name__, sequential_agent
)
events = [e async for e in sequential_agent.run_live(parent_ctx)]
assert len(events) == 2
assert events[0].author == agent_1.name
assert events[1].author == agent_2.name
assert events[0].content.parts[0].text == f'Hello, live {agent_1.name}!'
assert events[1].content.parts[0].text == f'Hello, live {agent_2.name}!'
src/google/adk/tests/unittests/artifacts/__init__.py
-14
View File
@@ -1,14 +0,0 @@
# 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.
src/google/adk/tests/unittests/artifacts/test_artifact_service.py
-276
View File
@@ -1,276 +0,0 @@
# 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 artifact service."""
import enum
from typing import Optional
from typing import Union
from google.adk.artifacts import GcsArtifactService
from google.adk.artifacts import InMemoryArtifactService
from google.genai import types
import pytest
Enum = enum.Enum
class ArtifactServiceType(Enum):
IN_MEMORY = "IN_MEMORY"
GCS = "GCS"
class MockBlob:
"""Mocks a GCS Blob object.
This class provides mock implementations for a few common GCS Blob methods,
allowing the user to test code that interacts with GCS without actually
connecting to a real bucket.
"""
def __init__(self, name: str) -> None:
"""Initializes a MockBlob.
Args:
name: The name of the blob.
"""
self.name = name
self.content: Optional[bytes] = None
self.content_type: Optional[str] = None
def upload_from_string(
self, data: Union[str, bytes], content_type: Optional[str] = None
) -> None:
"""Mocks uploading data to the blob (from a string or bytes).
Args:
data: The data to upload (string or bytes).
content_type: The content type of the data (optional).
"""
if isinstance(data, str):
self.content = data.encode("utf-8")
elif isinstance(data, bytes):
self.content = data
else:
raise TypeError("data must be str or bytes")
if content_type:
self.content_type = content_type
def download_as_bytes(self) -> bytes:
"""Mocks downloading the blob's content as bytes.
Returns:
bytes: The content of the blob as bytes.
Raises:
Exception: If the blob doesn't exist (hasn't been uploaded to).
"""
if self.content is None:
return b""
return self.content
def delete(self) -> None:
"""Mocks deleting a blob."""
self.content = None
self.content_type = None
class MockBucket:
"""Mocks a GCS Bucket object."""
def __init__(self, name: str) -> None:
"""Initializes a MockBucket.
Args:
name: The name of the bucket.
"""
self.name = name
self.blobs: dict[str, MockBlob] = {}
def blob(self, blob_name: str) -> MockBlob:
"""Mocks getting a Blob object (doesn't create it in storage).
Args:
blob_name: The name of the blob.
Returns:
A MockBlob instance.
"""
if blob_name not in self.blobs:
self.blobs[blob_name] = MockBlob(blob_name)
return self.blobs[blob_name]
class MockClient:
"""Mocks the GCS Client."""
def __init__(self) -> None:
"""Initializes MockClient."""
self.buckets: dict[str, MockBucket] = {}
def bucket(self, bucket_name: str) -> MockBucket:
"""Mocks getting a Bucket object."""
if bucket_name not in self.buckets:
self.buckets[bucket_name] = MockBucket(bucket_name)
return self.buckets[bucket_name]
def list_blobs(self, bucket: MockBucket, prefix: Optional[str] = None):
"""Mocks listing blobs in a bucket, optionally with a prefix."""
if prefix:
return [
blob for name, blob in bucket.blobs.items() if name.startswith(prefix)
]
return list(bucket.blobs.values())
def mock_gcs_artifact_service():
"""Creates a mock GCS artifact service for testing."""
service = GcsArtifactService(bucket_name="test_bucket")
service.storage_client = MockClient()
service.bucket = service.storage_client.bucket("test_bucket")
return service
def get_artifact_service(
service_type: ArtifactServiceType = ArtifactServiceType.IN_MEMORY,
):
"""Creates an artifact service for testing."""
if service_type == ArtifactServiceType.GCS:
return mock_gcs_artifact_service()
return InMemoryArtifactService()
@pytest.mark.parametrize(
"service_type", [ArtifactServiceType.IN_MEMORY, ArtifactServiceType.GCS]
)
def test_load_empty(service_type):
"""Tests loading an artifact when none exists."""
artifact_service = get_artifact_service(service_type)
assert not artifact_service.load_artifact(
app_name="test_app",
user_id="test_user",
session_id="session_id",
filename="filename",
)
@pytest.mark.parametrize(
"service_type", [ArtifactServiceType.IN_MEMORY, ArtifactServiceType.GCS]
)
def test_save_load_delete(service_type):
"""Tests saving, loading, and deleting an artifact."""
artifact_service = get_artifact_service(service_type)
artifact = types.Part.from_bytes(data=b"test_data", mime_type="text/plain")
app_name = "app0"
user_id = "user0"
session_id = "123"
filename = "file456"
artifact_service.save_artifact(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename=filename,
artifact=artifact,
)
assert (
artifact_service.load_artifact(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename=filename,
)
== artifact
)
artifact_service.delete_artifact(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename=filename,
)
assert not artifact_service.load_artifact(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename=filename,
)
@pytest.mark.parametrize(
"service_type", [ArtifactServiceType.IN_MEMORY, ArtifactServiceType.GCS]
)
def test_list_keys(service_type):
"""Tests listing keys in the artifact service."""
artifact_service = get_artifact_service(service_type)
artifact = types.Part.from_bytes(data=b"test_data", mime_type="text/plain")
app_name = "app0"
user_id = "user0"
session_id = "123"
filename = "filename"
filenames = [filename + str(i) for i in range(5)]
for f in filenames:
artifact_service.save_artifact(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename=f,
artifact=artifact,
)
assert (
artifact_service.list_artifact_keys(
app_name=app_name, user_id=user_id, session_id=session_id
)
== filenames
)
@pytest.mark.parametrize(
"service_type", [ArtifactServiceType.IN_MEMORY, ArtifactServiceType.GCS]
)
def test_list_versions(service_type):
"""Tests listing versions of an artifact."""
artifact_service = get_artifact_service(service_type)
app_name = "app0"
user_id = "user0"
session_id = "123"
filename = "filename"
versions = [
types.Part.from_bytes(
data=i.to_bytes(2, byteorder="big"), mime_type="text/plain"
)
for i in range(3)
]
for i in range(3):
artifact_service.save_artifact(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename=filename,
artifact=versions[i],
)
response_versions = artifact_service.list_versions(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename=filename,
)
assert response_versions == list(range(3))
src/google/adk/tests/unittests/auth/test_auth_handler.py
-575
View File
File diff suppressed because it is too large Load Diff
src/google/adk/tests/unittests/conftest.py
-73
View File
@@ -1,73 +0,0 @@
# 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
from pytest import fixture
from pytest import FixtureRequest
from pytest import hookimpl
from pytest import Metafunc
_ENV_VARS = {
'GOOGLE_API_KEY': 'fake_google_api_key',
'GOOGLE_CLOUD_PROJECT': 'fake_google_cloud_project',
'GOOGLE_CLOUD_LOCATION': 'fake_google_cloud_location',
}
ENV_SETUPS = {
'GOOGLE_AI': {
'GOOGLE_GENAI_USE_VERTEXAI': '0',
**_ENV_VARS,
},
'VERTEX': {
'GOOGLE_GENAI_USE_VERTEXAI': '1',
**_ENV_VARS,
},
}
@fixture(autouse=True)
def env_variables(request: FixtureRequest):
# Set up the environment
env_name: str = request.param
envs = ENV_SETUPS[env_name]
original_env = {key: os.environ.get(key) for key in envs}
os.environ.update(envs)
yield # Run the test
# Restore the environment
for key in envs:
if (original_val := original_env.get(key)) is None:
os.environ.pop(key, None)
else:
os.environ[key] = original_val
@hookimpl(tryfirst=True)
def pytest_generate_tests(metafunc: Metafunc):
"""Generate test cases for each environment setup."""
if env_variables.__name__ in metafunc.fixturenames:
if not _is_explicitly_marked(env_variables.__name__, metafunc):
metafunc.parametrize(
env_variables.__name__, ENV_SETUPS.keys(), indirect=True
)
def _is_explicitly_marked(mark_name: str, metafunc: Metafunc) -> bool:
if hasattr(metafunc.function, 'pytestmark'):
for mark in metafunc.function.pytestmark:
if mark.name == 'parametrize' and mark.args[0] == mark_name:
return True
return False
src/google/adk/tests/unittests/fast_api/__init__.py
-14
View File
@@ -1,14 +0,0 @@
# 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.
src/google/adk/tests/unittests/fast_api/test_fast_api.py
-269
View File
@@ -1,269 +0,0 @@
# 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 json
import sys
import threading
import time
import types as ptypes
from typing import AsyncGenerator
from google.adk.agents import BaseAgent
from google.adk.agents import LiveRequest
from google.adk.agents.run_config import RunConfig
from google.adk.cli.fast_api import AgentRunRequest
from google.adk.cli.fast_api import get_fast_api_app
from google.adk.cli.utils import envs
from google.adk.events import Event
from google.adk.runners import Runner
from google.genai import types
import httpx
import pytest
from uvicorn.main import run as uvicorn_run
import websockets
# Here we “fake” the agent module that get_fast_api_app expects.
# The server code does: `agent_module = importlib.import_module(agent_name)`
# and then accesses: agent_module.agent.root_agent.
class DummyAgent(BaseAgent):
pass
dummy_module = ptypes.ModuleType("test_agent")
dummy_module.agent = ptypes.SimpleNamespace(
root_agent=DummyAgent(name="dummy_agent")
)
sys.modules["test_app"] = dummy_module
envs.load_dotenv_for_agent("test_app", ".")
event1 = Event(
author="dummy agent",
invocation_id="invocation_id",
content=types.Content(
role="model", parts=[types.Part(text="LLM reply", inline_data=None)]
),
)
event2 = Event(
author="dummy agent",
invocation_id="invocation_id",
content=types.Content(
role="model",
parts=[
types.Part(
text=None,
inline_data=types.Blob(
mime_type="audio/pcm;rate=24000", data=b"\x00\xFF"
),
)
],
),
)
event3 = Event(
author="dummy agent", invocation_id="invocation_id", interrupted=True
)
# For simplicity, we patch Runner.run_live to yield dummy events.
# We use SimpleNamespace to mimic attribute-access (i.e. event.content.parts).
async def dummy_run_live(
self, session, live_request_queue
) -> AsyncGenerator[Event, None]:
# Immediately yield a dummy event with a text reply.
yield event1
await asyncio.sleep(0)
yield event2
await asyncio.sleep(0)
yield event3
raise Exception()
async def dummy_run_async(
self,
user_id,
session_id,
new_message,
run_config: RunConfig = RunConfig(),
) -> AsyncGenerator[Event, None]:
# Immediately yield a dummy event with a text reply.
yield event1
await asyncio.sleep(0)
yield event2
await asyncio.sleep(0)
yield event3
return
###############################################################################
# Pytest fixtures to patch methods and start the server
###############################################################################
@pytest.fixture(autouse=True)
def patch_runner(monkeypatch):
# Patch the Runner methods to use our dummy implementations.
monkeypatch.setattr(Runner, "run_live", dummy_run_live)
monkeypatch.setattr(Runner, "run_async", dummy_run_async)
@pytest.fixture(scope="module", autouse=True)
def start_server():
"""Start the FastAPI server in a background thread."""
def run_server():
uvicorn_run(
get_fast_api_app(agent_dir=".", web=True),
host="0.0.0.0",
log_config=None,
)
server_thread = threading.Thread(target=run_server, daemon=True)
server_thread.start()
# Wait a moment to ensure the server is up.
time.sleep(2)
yield
# The daemon thread will be terminated when tests complete.
@pytest.mark.asyncio
async def test_sse_endpoint():
base_http_url = "http://127.0.0.1:8000"
user_id = "test_user"
session_id = "test_session"
# Ensure that the session exists (create if necessary).
url_create = (
f"{base_http_url}/apps/test_app/users/{user_id}/sessions/{session_id}"
)
httpx.post(url_create, json={"state": {}})
async with httpx.AsyncClient() as client:
# Make a POST request to the SSE endpoint.
async with client.stream(
"POST",
f"{base_http_url}/run_sse",
json=json.loads(
AgentRunRequest(
app_name="test_app",
user_id=user_id,
session_id=session_id,
new_message=types.Content(
parts=[types.Part(text="Hello via SSE", inline_data=None)]
),
streaming=False,
).model_dump_json(exclude_none=True)
),
) as response:
# Ensure the status code and header are as expected.
assert response.status_code == 200
assert (
response.headers.get("content-type")
== "text/event-stream; charset=utf-8"
)
# Iterate over events from the stream.
event_count = 0
event_buffer = ""
async for line in response.aiter_lines():
event_buffer += line + "\n"
# An SSE event is terminated by an empty line (double newline)
if line == "" and event_buffer.strip():
# Process the complete event
event_data = None
for event_line in event_buffer.split("\n"):
if event_line.startswith("data: "):
event_data = event_line[6:] # Remove "data: " prefix
if event_data:
event_count += 1
if event_count == 1:
assert event_data == event1.model_dump_json(
exclude_none=True, by_alias=True
)
elif event_count == 2:
assert event_data == event2.model_dump_json(
exclude_none=True, by_alias=True
)
elif event_count == 3:
assert event_data == event3.model_dump_json(
exclude_none=True, by_alias=True
)
else:
pass
# Reset buffer for next event
event_buffer = ""
assert event_count == 3 # Expecting 3 events from dummy_run_async
@pytest.mark.asyncio
async def test_websocket_endpoint():
base_http_url = "http://127.0.0.1:8000"
base_ws_url = "ws://127.0.0.1:8000"
user_id = "test_user"
session_id = "test_session"
# Ensure that the session exists (create if necessary).
url_create = (
f"{base_http_url}/apps/test_app/users/{user_id}/sessions/{session_id}"
)
httpx.post(url_create, json={"state": {}})
ws_url = f"{base_ws_url}/run_live?app_name=test_app&user_id={user_id}&session_id={session_id}"
async with websockets.connect(ws_url) as ws:
# --- Test sending text data ---
text_payload = LiveRequest(
content=types.Content(
parts=[types.Part(text="Hello via WebSocket", inline_data=None)]
)
)
await ws.send(text_payload.model_dump_json())
# Wait for a reply from our dummy_run_live.
reply = await ws.recv()
event = Event.model_validate_json(reply)
assert event.content.parts[0].text == "LLM reply"
# --- Test sending binary data (allowed mime type "audio/pcm") ---
sample_audio = b"\x00\xFF"
binary_payload = LiveRequest(
blob=types.Blob(
mime_type="audio/pcm",
data=sample_audio,
)
)
await ws.send(binary_payload.model_dump_json())
# Wait for a reply.
reply = await ws.recv()
event = Event.model_validate_json(reply)
assert (
event.content.parts[0].inline_data.mime_type == "audio/pcm;rate=24000"
)
assert event.content.parts[0].inline_data.data == b"\x00\xFF"
reply = await ws.recv()
event = Event.model_validate_json(reply)
assert event.interrupted is True
assert event.content is None
src/google/adk/tests/unittests/flows/__init__.py
-14
View File
@@ -1,14 +0,0 @@
# 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.
src/google/adk/tests/unittests/flows/llm_flows/__init__.py
-14
View File
@@ -1,14 +0,0 @@
# 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.
src/google/adk/tests/unittests/flows/llm_flows/_test_examples.py
-142
View File
@@ -1,142 +0,0 @@
# 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.
# TODO: delete and rewrite unit tests
from google.adk.agents import Agent
from google.adk.examples import BaseExampleProvider
from google.adk.examples import Example
from google.adk.flows.llm_flows import examples
from google.adk.models.base_llm import LlmRequest
from google.genai import types
import pytest
from ... import utils
@pytest.mark.asyncio
async def test_no_examples():
request = LlmRequest(
model="gemini-1.5-flash",
config=types.GenerateContentConfig(system_instruction=""),
)
agent = Agent(model="gemini-1.5-flash", name="agent", examples=[])
invocation_context = utils.create_invocation_context(
agent=agent, user_content=""
)
async for _ in examples.request_processor.run_async(
invocation_context,
request,
):
pass
assert request.config.system_instruction == ""
@pytest.mark.asyncio
async def test_agent_examples():
example_list = [
Example(
input=types.Content(
role="user",
parts=[types.Part.from_text(text="test1")],
),
output=[
types.Content(
role="model",
parts=[types.Part.from_text(text="response1")],
),
],
)
]
request = LlmRequest(
model="gemini-1.5-flash",
config=types.GenerateContentConfig(system_instruction=""),
)
agent = Agent(
model="gemini-1.5-flash",
name="agent",
examples=example_list,
)
invocation_context = utils.create_invocation_context(
agent=agent, user_content="test"
)
async for _ in examples.request_processor.run_async(
invocation_context,
request,
):
pass
assert (
request.config.system_instruction
== "<EXAMPLES>\nBegin few-shot\nThe following are examples of user"
" queries and model responses using the available tools.\n\nEXAMPLE"
" 1:\nBegin example\n[user]\ntest1\n\n[model]\nresponse1\nEnd"
" example\n\nEnd few-shot\nNow, try to follow these examples and"
" complete the following conversation\n<EXAMPLES>"
)
@pytest.mark.asyncio
async def test_agent_base_example_provider():
class TestExampleProvider(BaseExampleProvider):
def get_examples(self, query: str) -> list[Example]:
if query == "test":
return [
Example(
input=types.Content(
role="user",
parts=[types.Part.from_text(text="test")],
),
output=[
types.Content(
role="model",
parts=[types.Part.from_text(text="response1")],
),
],
)
]
else:
return []
provider = TestExampleProvider()
request = LlmRequest(
model="gemini-1.5-flash",
config=types.GenerateContentConfig(system_instruction=""),
)
agent = Agent(
model="gemini-1.5-flash",
name="agent",
examples=provider,
)
invocation_context = utils.create_invocation_context(
agent=agent, user_content="test"
)
async for _ in examples.request_processor.run_async(
invocation_context,
request,
):
pass
assert (
request.config.system_instruction
== "<EXAMPLES>\nBegin few-shot\nThe following are examples of user"
" queries and model responses using the available tools.\n\nEXAMPLE"
" 1:\nBegin example\n[user]\ntest\n\n[model]\nresponse1\nEnd"
" example\n\nEnd few-shot\nNow, try to follow these examples and"
" complete the following conversation\n<EXAMPLES>"
)
src/google/adk/tests/unittests/flows/llm_flows/test_agent_transfer.py
-311
View File
@@ -1,311 +0,0 @@
# 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 google.adk.agents.llm_agent import Agent
from google.adk.agents.loop_agent import LoopAgent
from google.adk.agents.sequential_agent import SequentialAgent
from google.adk.tools import exit_loop
from google.genai.types import Part
from ... import utils
def transfer_call_part(agent_name: str) -> Part:
return Part.from_function_call(
name='transfer_to_agent', args={'agent_name': agent_name}
)
TRANSFER_RESPONSE_PART = Part.from_function_response(
name='transfer_to_agent', response={}
)
def test_auto_to_auto():
response = [
transfer_call_part('sub_agent_1'),
'response1',
'response2',
]
mockModel = utils.MockModel.create(responses=response)
# root (auto) - sub_agent_1 (auto)
sub_agent_1 = Agent(name='sub_agent_1', model=mockModel)
root_agent = Agent(
name='root_agent',
model=mockModel,
sub_agents=[sub_agent_1],
)
runner = utils.InMemoryRunner(root_agent)
# Asserts the transfer.
assert utils.simplify_events(runner.run('test1')) == [
('root_agent', transfer_call_part('sub_agent_1')),
('root_agent', TRANSFER_RESPONSE_PART),
('sub_agent_1', 'response1'),
]
# sub_agent_1 should still be the current agent.
assert utils.simplify_events(runner.run('test2')) == [
('sub_agent_1', 'response2'),
]
def test_auto_to_single():
response = [
transfer_call_part('sub_agent_1'),
'response1',
'response2',
]
mockModel = utils.MockModel.create(responses=response)
# root (auto) - sub_agent_1 (single)
sub_agent_1 = Agent(
name='sub_agent_1',
model=mockModel,
disallow_transfer_to_parent=True,
disallow_transfer_to_peers=True,
)
root_agent = Agent(
name='root_agent', model=mockModel, sub_agents=[sub_agent_1]
)
runner = utils.InMemoryRunner(root_agent)
# Asserts the responses.
assert utils.simplify_events(runner.run('test1')) == [
('root_agent', transfer_call_part('sub_agent_1')),
('root_agent', TRANSFER_RESPONSE_PART),
('sub_agent_1', 'response1'),
]
# root_agent should still be the current agent, becaues sub_agent_1 is single.
assert utils.simplify_events(runner.run('test2')) == [
('root_agent', 'response2'),
]
def test_auto_to_auto_to_single():
response = [
transfer_call_part('sub_agent_1'),
# sub_agent_1 transfers to sub_agent_1_1.
transfer_call_part('sub_agent_1_1'),
'response1',
'response2',
]
mockModel = utils.MockModel.create(responses=response)
# root (auto) - sub_agent_1 (auto) - sub_agent_1_1 (single)
sub_agent_1_1 = Agent(
name='sub_agent_1_1',
model=mockModel,
disallow_transfer_to_parent=True,
disallow_transfer_to_peers=True,
)
sub_agent_1 = Agent(
name='sub_agent_1', model=mockModel, sub_agents=[sub_agent_1_1]
)
root_agent = Agent(
name='root_agent', model=mockModel, sub_agents=[sub_agent_1]
)
runner = utils.InMemoryRunner(root_agent)
# Asserts the responses.
assert utils.simplify_events(runner.run('test1')) == [
('root_agent', transfer_call_part('sub_agent_1')),
('root_agent', TRANSFER_RESPONSE_PART),
('sub_agent_1', transfer_call_part('sub_agent_1_1')),
('sub_agent_1', TRANSFER_RESPONSE_PART),
('sub_agent_1_1', 'response1'),
]
# sub_agent_1 should still be the current agent. sub_agent_1_1 is single so it should
# not be the current agent, otherwise the conversation will be tied to
# sub_agent_1_1 forever.
assert utils.simplify_events(runner.run('test2')) == [
('sub_agent_1', 'response2'),
]
def test_auto_to_sequential():
response = [
transfer_call_part('sub_agent_1'),
# sub_agent_1 responds directly instead of transfering.
'response1',
'response2',
'response3',
]
mockModel = utils.MockModel.create(responses=response)
# root (auto) - sub_agent_1 (sequential) - sub_agent_1_1 (single)
# \ sub_agent_1_2 (single)
sub_agent_1_1 = Agent(
name='sub_agent_1_1',
model=mockModel,
disallow_transfer_to_parent=True,
disallow_transfer_to_peers=True,
)
sub_agent_1_2 = Agent(
name='sub_agent_1_2',
model=mockModel,
disallow_transfer_to_parent=True,
disallow_transfer_to_peers=True,
)
sub_agent_1 = SequentialAgent(
name='sub_agent_1',
sub_agents=[sub_agent_1_1, sub_agent_1_2],
)
root_agent = Agent(
name='root_agent',
model=mockModel,
sub_agents=[sub_agent_1],
)
runner = utils.InMemoryRunner(root_agent)
# Asserts the transfer.
assert utils.simplify_events(runner.run('test1')) == [
('root_agent', transfer_call_part('sub_agent_1')),
('root_agent', TRANSFER_RESPONSE_PART),
('sub_agent_1_1', 'response1'),
('sub_agent_1_2', 'response2'),
]
# root_agent should still be the current agent because sub_agent_1 is sequential.
assert utils.simplify_events(runner.run('test2')) == [
('root_agent', 'response3'),
]
def test_auto_to_sequential_to_auto():
response = [
transfer_call_part('sub_agent_1'),
# sub_agent_1 responds directly instead of transfering.
'response1',
transfer_call_part('sub_agent_1_2_1'),
'response2',
'response3',
'response4',
]
mockModel = utils.MockModel.create(responses=response)
# root (auto) - sub_agent_1 (seq) - sub_agent_1_1 (single)
# \ sub_agent_1_2 (auto) - sub_agent_1_2_1 (auto)
# \ sub_agent_1_3 (single)
sub_agent_1_1 = Agent(
name='sub_agent_1_1',
model=mockModel,
disallow_transfer_to_parent=True,
disallow_transfer_to_peers=True,
)
sub_agent_1_2_1 = Agent(name='sub_agent_1_2_1', model=mockModel)
sub_agent_1_2 = Agent(
name='sub_agent_1_2',
model=mockModel,
sub_agents=[sub_agent_1_2_1],
)
sub_agent_1_3 = Agent(
name='sub_agent_1_3',
model=mockModel,
disallow_transfer_to_parent=True,
disallow_transfer_to_peers=True,
)
sub_agent_1 = SequentialAgent(
name='sub_agent_1',
sub_agents=[sub_agent_1_1, sub_agent_1_2, sub_agent_1_3],
)
root_agent = Agent(
name='root_agent',
model=mockModel,
sub_agents=[sub_agent_1],
)
runner = utils.InMemoryRunner(root_agent)
# Asserts the transfer.
assert utils.simplify_events(runner.run('test1')) == [
('root_agent', transfer_call_part('sub_agent_1')),
('root_agent', TRANSFER_RESPONSE_PART),
('sub_agent_1_1', 'response1'),
('sub_agent_1_2', transfer_call_part('sub_agent_1_2_1')),
('sub_agent_1_2', TRANSFER_RESPONSE_PART),
('sub_agent_1_2_1', 'response2'),
('sub_agent_1_3', 'response3'),
]
# root_agent should still be the current agent because sub_agent_1 is sequential.
assert utils.simplify_events(runner.run('test2')) == [
('root_agent', 'response4'),
]
def test_auto_to_loop():
response = [
transfer_call_part('sub_agent_1'),
# sub_agent_1 responds directly instead of transfering.
'response1',
'response2',
'response3',
Part.from_function_call(name='exit_loop', args={}),
'response4',
'response5',
]
mockModel = utils.MockModel.create(responses=response)
# root (auto) - sub_agent_1 (loop) - sub_agent_1_1 (single)
# \ sub_agent_1_2 (single)
sub_agent_1_1 = Agent(
name='sub_agent_1_1',
model=mockModel,
disallow_transfer_to_parent=True,
disallow_transfer_to_peers=True,
)
sub_agent_1_2 = Agent(
name='sub_agent_1_2',
model=mockModel,
disallow_transfer_to_parent=True,
disallow_transfer_to_peers=True,
tools=[exit_loop],
)
sub_agent_1 = LoopAgent(
name='sub_agent_1',
sub_agents=[sub_agent_1_1, sub_agent_1_2],
)
root_agent = Agent(
name='root_agent',
model=mockModel,
sub_agents=[sub_agent_1],
)
runner = utils.InMemoryRunner(root_agent)
# Asserts the transfer.
assert utils.simplify_events(runner.run('test1')) == [
# Transfers to sub_agent_1.
('root_agent', transfer_call_part('sub_agent_1')),
('root_agent', TRANSFER_RESPONSE_PART),
# Loops.
('sub_agent_1_1', 'response1'),
('sub_agent_1_2', 'response2'),
('sub_agent_1_1', 'response3'),
# Exits.
('sub_agent_1_2', Part.from_function_call(name='exit_loop', args={})),
(
'sub_agent_1_2',
Part.from_function_response(name='exit_loop', response={}),
),
# root_agent summarizes.
('root_agent', 'response4'),
]
# root_agent should still be the current agent because sub_agent_1 is loop.
assert utils.simplify_events(runner.run('test2')) == [
('root_agent', 'response5'),
]

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