Files
google--adk-python/tests/unittests/flows/llm_flows/test_tool_telemetry.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

100 lines
3.2 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.
from typing import Any
from typing import Dict
from typing import Optional
from unittest import mock
from google.adk.agents.llm_agent import Agent
from google.adk.events.event import Event
from google.adk.flows.llm_flows.functions import handle_function_calls_async
from google.adk.telemetry import tracing
from google.adk.tools.function_tool import FunctionTool
from google.genai import types
from ... import testing_utils
async def invoke_tool() -> Optional[Event]:
def simple_fn(**kwargs) -> Dict[str, Any]:
return {'result': 'test'}
tool = FunctionTool(simple_fn)
model = testing_utils.MockModel.create(responses=[])
agent = Agent(
name='agent',
model=model,
tools=[tool],
)
invocation_context = await testing_utils.create_invocation_context(
agent=agent, user_content=''
)
function_call = types.FunctionCall(name=tool.name, args={'a': 1, 'b': 2})
content = types.Content(parts=[types.Part(function_call=function_call)])
event = Event(
invocation_id=invocation_context.invocation_id,
author=agent.name,
content=content,
)
tools_dict = {tool.name: tool}
return await handle_function_calls_async(
invocation_context,
event,
tools_dict,
)
async def test_simple_function_with_mocked_tracer(monkeypatch):
mock_start_as_current_span_func = mock.Mock()
returned_context_manager_mock = mock.MagicMock()
returned_context_manager_mock.__enter__.return_value = mock.Mock(
name='span_mock'
)
mock_start_as_current_span_func.return_value = returned_context_manager_mock
monkeypatch.setattr(
tracing.tracer, 'start_as_current_span', mock_start_as_current_span_func
)
mock_adk_trace_tool_call = mock.Mock()
monkeypatch.setattr(
'google.adk.telemetry.tracing.trace_tool_call',
mock_adk_trace_tool_call,
)
event = await invoke_tool()
assert event is not None
event = await invoke_tool()
assert event is not None
expected_span_name = 'execute_tool simple_fn'
assert mock_start_as_current_span_func.call_count == 2
mock_start_as_current_span_func.assert_any_call(expected_span_name)
assert returned_context_manager_mock.__enter__.call_count == 2
assert returned_context_manager_mock.__exit__.call_count == 2
assert mock_adk_trace_tool_call.call_count == 2
for call_args_item in mock_adk_trace_tool_call.call_args_list:
kwargs = call_args_item.kwargs
assert kwargs['tool'].name == 'simple_fn'
assert kwargs['args'] == {'a': 1, 'b': 2}
assert 'function_response_event' in kwargs
assert kwargs['function_response_event'].content.parts[
0
].function_response.response == {'result': 'test'}