diff --git a/tests/unittests/telemetry/__init__.py b/tests/unittests/telemetry/__init__.py new file mode 100644 index 0000000000..0a2669d7a2 --- /dev/null +++ b/tests/unittests/telemetry/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/tests/unittests/telemetry/test_functional.py b/tests/unittests/telemetry/test_functional.py new file mode 100644 index 0000000000..54db282a6c --- /dev/null +++ b/tests/unittests/telemetry/test_functional.py @@ -0,0 +1,129 @@ +# 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 gc +import sys +from unittest import mock + +from google.adk import telemetry +from google.adk.agents import base_agent +from google.adk.agents.llm_agent import Agent +from google.adk.models.base_llm import BaseLlm +from google.adk.tools import FunctionTool +from google.adk.utils.context_utils import Aclosing +from google.genai.types import Part +from opentelemetry.version import __version__ +import pytest + +from ..testing_utils import MockModel +from ..testing_utils import TestInMemoryRunner + + +@pytest.fixture +def test_model() -> BaseLlm: + mock_model = MockModel.create( + responses=[ + Part.from_function_call(name='some_tool', args={}), + Part.from_text(text='text response'), + ] + ) + return mock_model + + +@pytest.fixture +def test_agent(test_model: BaseLlm) -> Agent: + def some_tool(): + pass + + root_agent = Agent( + name='some_root_agent', + model=test_model, + tools=[ + FunctionTool(some_tool), + ], + ) + return root_agent + + +@pytest.fixture +async def test_runner(test_agent: Agent) -> TestInMemoryRunner: + runner = TestInMemoryRunner(test_agent) + return runner + + +@pytest.fixture +def mock_start_as_current_span(monkeypatch: pytest.MonkeyPatch) -> mock.Mock: + mock_context_manager = mock.MagicMock() + mock_context_manager.__enter__.return_value = mock.Mock() + mock_start_as_current_span = mock.Mock() + mock_start_as_current_span.return_value = mock_context_manager + + def do_replace(tracer): + monkeypatch.setattr( + tracer, 'start_as_current_span', mock_start_as_current_span + ) + + do_replace(telemetry.tracer) + do_replace(base_agent.tracer) + + return mock_start_as_current_span + + +@pytest.mark.asyncio +async def test_tracer_start_as_current_span( + test_runner: TestInMemoryRunner, + mock_start_as_current_span: mock.Mock, +): + """Test creation of multiple spans in an E2E runner invocation. + + Additionally tests if each async generator invoked is wrapped in Aclosing. + This is necessary because instrumentation utilizes contextvars, which ran into "ContextVar was created in a different Context" errors, + when a given coroutine gets indeterminitely suspended. + """ + firstiter, finalizer = sys.get_asyncgen_hooks() + + def wrapped_firstiter(coro): + nonlocal firstiter + assert any( + isinstance(referrer, Aclosing) + or isinstance(indirect_referrer, Aclosing) + for referrer in gc.get_referrers(coro) + # Some coroutines have a layer of indirection in python 3.9 and 3.10 + for indirect_referrer in gc.get_referrers(referrer) + ), f'Coro `{coro.__name__}` is not wrapped with Aclosing' + firstiter(coro) + + sys.set_asyncgen_hooks(wrapped_firstiter, finalizer) + + # Act + async with Aclosing(test_runner.run_async_with_new_session_agen('')) as agen: + async for _ in agen: + pass + + # Assert + expected_start_as_current_span_calls = [ + mock.call('invocation'), + mock.call('execute_tool some_tool'), + mock.call('agent_run [some_root_agent]'), + mock.call('call_llm'), + mock.call('call_llm'), + ] + + mock_start_as_current_span.assert_has_calls( + expected_start_as_current_span_calls, + any_order=True, + ) + assert mock_start_as_current_span.call_count == len( + expected_start_as_current_span_calls + ) diff --git a/tests/unittests/test_telemetry.py b/tests/unittests/telemetry/test_spans.py similarity index 100% rename from tests/unittests/test_telemetry.py rename to tests/unittests/telemetry/test_spans.py diff --git a/tests/unittests/testing_utils.py b/tests/unittests/testing_utils.py index 1ff3d79866..0cc637871b 100644 --- a/tests/unittests/testing_utils.py +++ b/tests/unittests/testing_utils.py @@ -36,6 +36,7 @@ from google.adk.runners import Runner from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.adk.sessions.session import Session +from google.adk.utils.context_utils import Aclosing from google.genai import types from google.genai.types import Part from typing_extensions import override @@ -161,19 +162,26 @@ async def run_async_with_new_session( self, new_message: types.ContentUnion ) -> list[Event]: + collected_events: list[Event] = [] + async for event in self.run_async_with_new_session_agen(new_message): + collected_events.append(event) + + return collected_events + + async def run_async_with_new_session_agen( + self, new_message: types.ContentUnion + ) -> AsyncGenerator[Event, None]: session = await self.session_service.create_session( app_name='InMemoryRunner', user_id='test_user' ) - collected_events = [] - - async for event in self.run_async( + agen = self.run_async( user_id=session.user_id, session_id=session.id, new_message=get_user_content(new_message), - ): - collected_events.append(event) - - return collected_events + ) + async with Aclosing(agen): + async for event in agen: + yield event class InMemoryRunner: