From d8476e87ad69fa1f15af886ecc90a908b627e412 Mon Sep 17 00:00:00 2001 From: Alexander Bondarenko Date: Fri, 21 Nov 2025 12:05:37 +0100 Subject: [PATCH 1/5] feat(testing): Add public testing utilities module - Add google.adk.testing module with MockModel, InMemoryRunner, and test helpers - Export testing utilities from src/google/adk/testing/ for PyPI distribution - Update tests/unittests/testing_utils.py to re-export from new public location - Add comprehensive README with examples and best practices - Maintain backward compatibility for existing ADK tests This makes testing utilities officially supported and available to ADK users via the PyPI package, enabling easier unit testing of custom agents without requiring real LLM API calls. --- src/google/adk/testing/README.md | 245 ++++++++++ src/google/adk/testing/__init__.py | 90 ++++ src/google/adk/testing/testing_utils.py | 619 ++++++++++++++++++++++++ tests/unittests/testing_utils.py | 442 ++--------------- 4 files changed, 995 insertions(+), 401 deletions(-) create mode 100644 src/google/adk/testing/README.md create mode 100644 src/google/adk/testing/__init__.py create mode 100644 src/google/adk/testing/testing_utils.py diff --git a/src/google/adk/testing/README.md b/src/google/adk/testing/README.md new file mode 100644 index 0000000000..7456ec2b71 --- /dev/null +++ b/src/google/adk/testing/README.md @@ -0,0 +1,245 @@ +# ADK Testing Utilities + +This module provides testing utilities for developers building agents with the Agent Development Kit (ADK). These tools make it easy to write unit tests for your agents without requiring real LLM API calls. + +## Quick Start + +```python +from google.adk import Agent +from google.adk.testing import MockModel, InMemoryRunner + +# Create an agent with a mock model +agent = Agent( + name="test_agent", + model=MockModel.create(responses=["Hello, I'm a test response!"]), + instruction="You are a helpful assistant." +) + +# Run the agent in a test environment +runner = InMemoryRunner(root_agent=agent) +events = runner.run("Hi there!") + +# Assert on the results +assert len(events) > 0 +``` + +## Key Components + +### MockModel + +A mock LLM that returns pre-defined responses instead of calling a real API. This makes tests fast, deterministic, and doesn't require API keys. + +```python +from google.adk.testing import MockModel + +# Simple text responses +mock = MockModel.create(responses=["Response 1", "Response 2"]) + +# Multiple responses for a conversation +mock = MockModel.create(responses=[ + "First response", + "Second response", + "Third response" +]) + +# Mock an error +mock = MockModel.create(responses=[], error=ValueError("API Error")) +``` + +### InMemoryRunner + +A test runner that uses in-memory services for fast, isolated tests. + +```python +from google.adk.testing import InMemoryRunner + +runner = InMemoryRunner(root_agent=agent) + +# Synchronous execution +events = runner.run("Test message") + +# Async execution +events = await runner.run_async("Test message") + +# Access the session +session = runner.session +``` + +### Helper Functions + +#### simplify_events() + +Simplify events for easier assertions in tests: + +```python +from google.adk.testing import simplify_events + +events = runner.run("Hello") +simplified = simplify_events(events) + +# Returns list of (author, simplified_content) tuples +# Easier to assert: [("user", "Hello"), ("test_agent", "Hi there!")] +``` + +#### create_test_agent() + +Create a basic test agent quickly: + +```python +from google.adk.testing import create_test_agent + +agent = create_test_agent(name="my_test") +``` + +#### create_invocation_context() + +Create a test invocation context for testing agent components: + +```python +from google.adk.testing import create_invocation_context + +context = await create_invocation_context( + agent=agent, + user_content="Test message", + run_config=RunConfig(), + plugins=[] +) +``` + +## Complete Example + +```python +import pytest +from google.adk import Agent +from google.adk.testing import MockModel, InMemoryRunner, simplify_events + +@pytest.fixture +def mock_agent(): + """Create a test agent with mock responses.""" + mock_model = MockModel.create(responses=[ + "I can help with that!", + "Here's the information you requested.", + "Is there anything else?" + ]) + + return Agent( + name="test_assistant", + model=mock_model, + instruction="You are a helpful assistant." + ) + +def test_basic_conversation(mock_agent): + """Test a basic conversation flow.""" + runner = InMemoryRunner(root_agent=mock_agent) + + # First message + events = runner.run("Can you help me?") + simplified = simplify_events(events) + + assert len(simplified) >= 2 + assert simplified[0][0] == "user" + assert simplified[0][1] == "Can you help me?" + + # Second message in same session + events = runner.run("Thanks!") + simplified = simplify_events(events) + + assert len(simplified) >= 2 + +@pytest.mark.asyncio +async def test_async_execution(mock_agent): + """Test async agent execution.""" + runner = InMemoryRunner(root_agent=mock_agent) + + events = await runner.run_async("Hello!") + assert len(events) > 0 + +def test_error_handling(): + """Test error handling with mock models.""" + error_model = MockModel.create( + responses=[], + error=ValueError("API unavailable") + ) + + agent = Agent(name="error_agent", model=error_model) + runner = InMemoryRunner(root_agent=agent) + + with pytest.raises(ValueError, match="API unavailable"): + runner.run("This will fail") +``` + +## API Reference + +### MockModel + +- `MockModel.create(responses, error=None)` - Create a mock model with pre-defined responses +- `model.requests` - List of all LlmRequest objects sent to the mock +- `model.responses` - List of LlmResponse objects the mock will return +- `model.response_index` - Current position in the responses list + +### InMemoryRunner + +- `InMemoryRunner(root_agent, plugins=[], app=None)` - Create a test runner +- `runner.run(message)` - Run agent synchronously, returns list of Events +- `runner.run_async(message, invocation_id=None)` - Run agent asynchronously +- `runner.session` - Access the current test session + +### Helper Functions + +- `create_test_agent(name)` - Create a simple test agent +- `create_invocation_context(agent, user_content, run_config, plugins)` - Create test context +- `simplify_events(events)` - Simplify events for assertions +- `simplify_content(content)` - Simplify content for assertions +- `append_user_content(context, parts)` - Add user content to context + +## Best Practices + +1. **Use MockModel for unit tests** - Fast, deterministic, no API calls needed +2. **Use InMemoryRunner** - Isolated test environment with in-memory services +3. **Use simplify_events() for assertions** - Makes test assertions cleaner +4. **Test both sync and async paths** - If your agent supports both +5. **Test error handling** - Use MockModel.create() with error parameter +6. **Keep tests isolated** - Each test should use its own runner instance + +## Integration Testing + +For integration tests with real LLM APIs, use the regular `Runner` instead: + +```python +from google.adk import Agent, Runner +from google.adk.sessions import InMemorySessionService + +# Real agent for integration testing +agent = Agent( + name="real_agent", + model="gemini-2.0-flash-exp", + instruction="You are a helpful assistant." +) + +runner = Runner( + agent=agent, + session_service=InMemorySessionService() +) + +# This will make real API calls +events = list(runner.run( + user_id="test_user", + session_id="test_session", + new_message="Hello!" +)) +``` + +## Migration Guide + +If you were previously using internal testing utilities from `tests/unittests/testing_utils.py`, you can now import from the public module: + +```python +# Old (internal) +from tests.unittests.testing_utils import MockModel, InMemoryRunner + +# New (public) +from google.adk.testing import MockModel, InMemoryRunner +``` + +All functionality remains the same - this just makes the utilities officially supported and available via the PyPI package. + diff --git a/src/google/adk/testing/__init__.py b/src/google/adk/testing/__init__.py new file mode 100644 index 0000000000..6ceca0c46a --- /dev/null +++ b/src/google/adk/testing/__init__.py @@ -0,0 +1,90 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Testing utilities for ADK users. + +This module provides utilities for testing agents built with the Agent +Development Kit (ADK). It includes mock models, test runners, and helper +functions to make testing agents easier without requiring real LLM API calls. + +Example: + Basic usage with MockModel: + + >>> from google.adk.testing import MockModel, InMemoryRunner, create_test_agent + >>> from google.adk import Agent + >>> + >>> # Create a test agent with a mock model + >>> agent = Agent( + ... name="test_agent", + ... model=MockModel.create(responses=["Hello, I'm a test response!"]), + ... instruction="You are a helpful assistant." + ... ) + >>> + >>> # Run the agent in a test environment + >>> runner = InMemoryRunner(root_agent=agent) + >>> events = runner.run("Hi there!") + >>> assert len(events) > 0 + + Testing with multiple responses: + + >>> mock_model = MockModel.create(responses=[ + ... "First response", + ... "Second response", + ... "Third response" + ... ]) + >>> # Each call to the model will return the next response in order + + Using helper functions for assertions: + + >>> from google.adk.testing import simplify_events + >>> events = runner.run("Test message") + >>> simplified = simplify_events(events) + >>> # Makes it easier to assert on event content +""" + +from __future__ import annotations + +from .testing_utils import append_user_content +from .testing_utils import create_invocation_context +from .testing_utils import create_test_agent +from .testing_utils import END_OF_AGENT +from .testing_utils import get_user_content +from .testing_utils import InMemoryRunner +from .testing_utils import MockLlmConnection +from .testing_utils import MockModel +from .testing_utils import ModelContent +from .testing_utils import simplify_content +from .testing_utils import simplify_contents +from .testing_utils import simplify_events +from .testing_utils import simplify_resumable_app_events +from .testing_utils import TestInMemoryRunner +from .testing_utils import UserContent + +__all__ = [ + 'MockModel', + 'MockLlmConnection', + 'InMemoryRunner', + 'TestInMemoryRunner', + 'create_test_agent', + 'create_invocation_context', + 'append_user_content', + 'simplify_events', + 'simplify_resumable_app_events', + 'simplify_contents', + 'simplify_content', + 'get_user_content', + 'UserContent', + 'ModelContent', + 'END_OF_AGENT', +] diff --git a/src/google/adk/testing/testing_utils.py b/src/google/adk/testing/testing_utils.py new file mode 100644 index 0000000000..758bb499c0 --- /dev/null +++ b/src/google/adk/testing/testing_utils.py @@ -0,0 +1,619 @@ +# 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. + +"""Public testing utilities for ADK users. + +This module provides utilities for testing agents built with the Agent +Development Kit (ADK). It includes mock models, test runners, and helper +functions to make testing agents easier. + +Example: + >>> from google.adk.testing import MockModel, InMemoryRunner, create_test_agent + >>> agent = create_test_agent(name="my_test_agent") + >>> mock_model = MockModel.create(responses=["Hello, world!"]) + >>> agent.model = mock_model + >>> runner = InMemoryRunner(root_agent=agent) + >>> events = runner.run("Hi there!") +""" + +from __future__ import annotations + +import asyncio +import contextlib +from typing import AsyncGenerator +from typing import Generator +from typing import Optional +from typing import Union + +from google.adk.agents.invocation_context import InvocationContext +from google.adk.agents.live_request_queue import LiveRequestQueue +from google.adk.agents.llm_agent import Agent +from google.adk.agents.llm_agent import LlmAgent +from google.adk.agents.run_config import RunConfig +from google.adk.apps.app import App +from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService +from google.adk.events.event import Event +from google.adk.memory.in_memory_memory_service import InMemoryMemoryService +from google.adk.models.base_llm import BaseLlm +from google.adk.models.base_llm_connection import BaseLlmConnection +from google.adk.models.llm_request import LlmRequest +from google.adk.models.llm_response import LlmResponse +from google.adk.plugins.base_plugin import BasePlugin +from google.adk.plugins.plugin_manager import PluginManager +from google.adk.runners import InMemoryRunner as AfInMemoryRunner +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 + + +def create_test_agent(name: str = 'test_agent') -> LlmAgent: + """Create a simple test agent for use in unit tests. + + Args: + name: The name of the test agent. + + Returns: + A configured LlmAgent instance suitable for testing. + """ + return LlmAgent(name=name) + + +class UserContent(types.Content): + """Helper to create user content for testing.""" + + def __init__(self, text_or_part: str): + parts = [ + types.Part.from_text(text=text_or_part) + if isinstance(text_or_part, str) + else text_or_part + ] + super().__init__(role='user', parts=parts) + + +class ModelContent(types.Content): + """Helper to create model content for testing.""" + + def __init__(self, parts: list[types.Part]): + super().__init__(role='model', parts=parts) + + +async def create_invocation_context( + agent: Agent, + user_content: str = '', + run_config: RunConfig = None, + plugins: list[BasePlugin] = [], +): + """Create an invocation context for testing agent components. + + Args: + agent: The agent to create context for. + user_content: Optional user message content. + run_config: Optional run configuration. + plugins: Optional list of plugins to use. + + Returns: + An InvocationContext configured for testing. + """ + invocation_id = 'test_id' + artifact_service = InMemoryArtifactService() + session_service = InMemorySessionService() + memory_service = InMemoryMemoryService() + invocation_context = InvocationContext( + artifact_service=artifact_service, + session_service=session_service, + memory_service=memory_service, + plugin_manager=PluginManager(plugins=plugins), + invocation_id=invocation_id, + agent=agent, + session=await session_service.create_session( + app_name='test_app', user_id='test_user' + ), + user_content=types.Content( + role='user', parts=[types.Part.from_text(text=user_content)] + ), + run_config=run_config or RunConfig(), + ) + if user_content: + append_user_content( + invocation_context, [types.Part.from_text(text=user_content)] + ) + return invocation_context + + +def append_user_content( + invocation_context: InvocationContext, parts: list[types.Part] +) -> Event: + """Append user content to an invocation context's session. + + Args: + invocation_context: The context to append to. + parts: Content parts to append. + + Returns: + The created Event. + """ + session = invocation_context.session + event = Event( + invocation_id=invocation_context.invocation_id, + author='user', + content=types.Content(role='user', parts=parts), + ) + session.events.append(event) + return event + + +# Extracts the contents from the events and transform them into a list of +# (author, simplified_content) tuples. +def simplify_events(events: list[Event]) -> list[(str, types.Part)]: + """Simplify events for easier assertion in tests. + + Args: + events: List of events to simplify. + + Returns: + List of (author, simplified_content) tuples. + """ + return [ + (event.author, simplify_content(event.content)) + for event in events + if event.content + ] + + +END_OF_AGENT = 'end_of_agent' + + +# Extracts the contents from the events and transform them into a list of +# (author, simplified_content OR AgentState OR "end_of_agent") tuples. +# +# Could be used to compare events for testing resumability. +def simplify_resumable_app_events( + events: list[Event], +) -> list[(str, Union[types.Part, str])]: + """Simplify events including agent state for resumability testing. + + Args: + events: List of events to simplify. + + Returns: + List of (author, simplified_content/state) tuples. + """ + results = [] + for event in events: + if event.content: + results.append((event.author, simplify_content(event.content))) + elif event.actions.end_of_agent: + results.append((event.author, END_OF_AGENT)) + elif event.actions.agent_state is not None: + results.append((event.author, event.actions.agent_state)) + return results + + +# Simplifies the contents into a list of (author, simplified_content) tuples. +def simplify_contents(contents: list[types.Content]) -> list[(str, types.Part)]: + """Simplify contents for easier assertion in tests. + + Args: + contents: List of contents to simplify. + + Returns: + List of (role, simplified_content) tuples. + """ + return [(content.role, simplify_content(content)) for content in contents] + + +# Simplifies the content so it's easier to assert. +# - If there is only one part, return part +# - If the only part is pure text, return stripped_text +# - If there are multiple parts, return parts +# - remove function_call_id if it exists +def simplify_content( + content: types.Content, +) -> Union[str, types.Part, list[types.Part]]: + """Simplify content for easier assertion in tests. + + Removes function call IDs and returns simplified representation: + - Single text part: returns stripped text string + - Single non-text part: returns the part + - Multiple parts: returns list of parts + + Args: + content: Content to simplify. + + Returns: + Simplified content representation. + """ + for part in content.parts: + if part.function_call and part.function_call.id: + part.function_call.id = None + if part.function_response and part.function_response.id: + part.function_response.id = None + if len(content.parts) == 1: + if content.parts[0].text: + return content.parts[0].text.strip() + else: + return content.parts[0] + return content.parts + + +def get_user_content(message: types.ContentUnion) -> types.Content: + """Convert a message to user Content. + + Args: + message: Either a Content object or string message. + + Returns: + A Content object with user role. + """ + return message if isinstance(message, types.Content) else UserContent(message) + + +class TestInMemoryRunner(AfInMemoryRunner): + """InMemoryRunner tailored for async tests. + + This runner extends the base InMemoryRunner with async run methods + suitable for unit testing. + """ + + async def run_async_with_new_session( + self, new_message: types.ContentUnion + ) -> list[Event]: + """Run agent with a new session and collect all events. + + Args: + new_message: The user message to send. + + Returns: + List of all events generated during execution. + """ + 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]: + """Run agent with a new session as an async generator. + + Args: + new_message: The user message to send. + + Yields: + Events as they are generated. + """ + session = await self.session_service.create_session( + app_name='InMemoryRunner', user_id='test_user' + ) + agen = self.run_async( + user_id=session.user_id, + session_id=session.id, + new_message=get_user_content(new_message), + ) + async with Aclosing(agen): + async for event in agen: + yield event + + +class InMemoryRunner: + """Test runner with in-memory services for testing agents. + + This runner is designed specifically for testing, providing easy access + to session state and in-memory services. + + Example: + >>> agent = create_test_agent() + >>> runner = InMemoryRunner(root_agent=agent) + >>> events = runner.run("Hello!") + >>> assert len(events) > 0 + """ + + def __init__( + self, + root_agent: Optional[Union[Agent, LlmAgent]] = None, + response_modalities: list[str] = None, + plugins: list[BasePlugin] = [], + app: Optional[App] = None, + ): + """Initializes the InMemoryRunner. + + Args: + root_agent: The root agent to run, won't be used if app is provided. + response_modalities: The response modalities of the runner. + plugins: The plugins to use in the runner, won't be used if app is + provided. + app: The app to use in the runner. + """ + if not app: + self.app_name = 'test_app' + self.root_agent = root_agent + self.runner = Runner( + app_name='test_app', + agent=root_agent, + artifact_service=InMemoryArtifactService(), + session_service=InMemorySessionService(), + memory_service=InMemoryMemoryService(), + plugins=plugins, + ) + else: + self.app_name = app.name + self.root_agent = app.root_agent + self.runner = Runner( + app=app, + artifact_service=InMemoryArtifactService(), + session_service=InMemorySessionService(), + memory_service=InMemoryMemoryService(), + ) + self.session_id = None + + @property + def session(self) -> Session: + """Get or create the test session.""" + if not self.session_id: + session = self.runner.session_service.create_session_sync( + app_name=self.app_name, user_id='test_user' + ) + self.session_id = session.id + return session + return self.runner.session_service.get_session_sync( + app_name=self.app_name, user_id='test_user', session_id=self.session_id + ) + + def run(self, new_message: types.ContentUnion) -> list[Event]: + """Run the agent synchronously and return all events. + + Args: + new_message: The user message to send. + + Returns: + List of all events generated during execution. + """ + return list( + self.runner.run( + user_id=self.session.user_id, + session_id=self.session.id, + new_message=get_user_content(new_message), + ) + ) + + async def run_async( + self, + new_message: Optional[types.ContentUnion] = None, + invocation_id: Optional[str] = None, + ) -> list[Event]: + """Run the agent asynchronously and return all events. + + Args: + new_message: The user message to send. + invocation_id: Optional invocation ID for tracking. + + Returns: + List of all events generated during execution. + """ + events = [] + async for event in self.runner.run_async( + user_id=self.session.user_id, + session_id=self.session.id, + invocation_id=invocation_id, + new_message=get_user_content(new_message) if new_message else None, + ): + events.append(event) + return events + + def run_live( + self, live_request_queue: LiveRequestQueue, run_config: RunConfig = None + ) -> list[Event]: + """Run the agent in live mode (bi-directional streaming). + + Args: + live_request_queue: Queue for live requests. + run_config: Optional run configuration. + + Returns: + List of events from the live session. + """ + collected_responses = [] + + async def consume_responses(session: Session): + run_res = self.runner.run_live( + session=session, + live_request_queue=live_request_queue, + run_config=run_config or RunConfig(), + ) + + async for response in run_res: + collected_responses.append(response) + # When we have enough response, we should return + if len(collected_responses) >= 1: + return + + try: + session = self.session + asyncio.run(consume_responses(session)) + except asyncio.TimeoutError: + print('Returning any partial results collected so far.') + + return collected_responses + + +class MockModel(BaseLlm): + """Mock LLM model for testing without real API calls. + + This model returns pre-defined responses instead of calling a real LLM API, + making tests fast, deterministic, and not requiring API keys. + + Example: + >>> mock = MockModel.create(responses=["Response 1", "Response 2"]) + >>> agent = LlmAgent(name="test", model=mock) + >>> # Agent will return "Response 1" on first call, "Response 2" on second + """ + + model: str = 'mock' + + requests: list[LlmRequest] = [] + responses: list[LlmResponse] + error: Union[Exception, None] = None + response_index: int = -1 + + @classmethod + def create( + cls, + responses: Union[ + list[types.Part], list[LlmResponse], list[str], list[list[types.Part]] + ], + error: Union[Exception, None] = None, + ): + """Create a MockModel with pre-defined responses. + + Args: + responses: List of responses to return. Can be: + - list[str]: Simple text responses + - list[Part]: Content parts + - list[list[Part]]: Multi-part responses + - list[LlmResponse]: Full response objects + error: Optional exception to raise instead of returning responses. + + Returns: + A configured MockModel instance. + """ + if error and not responses: + return cls(responses=[], error=error) + if not responses: + return cls(responses=[]) + elif isinstance(responses[0], LlmResponse): + # responses is list[LlmResponse] + return cls(responses=responses) + else: + responses = [ + LlmResponse(content=ModelContent(item)) + if isinstance(item, list) and isinstance(item[0], types.Part) + # responses is list[list[Part]] + else LlmResponse( + content=ModelContent( + # responses is list[str] or list[Part] + [Part(text=item) if isinstance(item, str) else item] + ) + ) + for item in responses + if item + ] + + return cls(responses=responses) + + @classmethod + @override + def supported_models(cls) -> list[str]: + """Return list of supported model names.""" + return ['mock'] + + def generate_content( + self, llm_request: LlmRequest, stream: bool = False + ) -> Generator[LlmResponse, None, None]: + """Generate content synchronously. + + Args: + llm_request: The request to process. + stream: Whether to stream the response (ignored for mock). + + Yields: + The next pre-defined response. + + Raises: + Exception: If an error was configured. + """ + if self.error is not None: + raise self.error + # Increasement of the index has to happen before the yield. + self.response_index += 1 + self.requests.append(llm_request) + # yield LlmResponse(content=self.responses[self.response_index]) + yield self.responses[self.response_index] + + @override + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + """Generate content asynchronously. + + Args: + llm_request: The request to process. + stream: Whether to stream the response (ignored for mock). + + Yields: + The next pre-defined response. + + Raises: + Exception: If an error was configured. + """ + if self.error is not None: + raise self.error + # Increasement of the index has to happen before the yield. + self.response_index += 1 + self.requests.append(llm_request) + yield self.responses[self.response_index] + + @contextlib.asynccontextmanager + async def connect(self, llm_request: LlmRequest) -> BaseLlmConnection: + """Create a live connection to the mock LLM. + + Args: + llm_request: The request to process. + + Yields: + A mock connection that returns pre-defined responses. + """ + self.requests.append(llm_request) + yield MockLlmConnection(self.responses) + + +class MockLlmConnection(BaseLlmConnection): + """Mock LLM connection for live/streaming tests.""" + + def __init__(self, llm_responses: list[LlmResponse]): + """Initialize with pre-defined responses. + + Args: + llm_responses: Responses to return. + """ + self.llm_responses = llm_responses + + async def send_history(self, history: list[types.Content]): + """Send history (no-op for mock).""" + pass + + async def send_content(self, content: types.Content): + """Send content (no-op for mock).""" + pass + + async def send(self, data): + """Send data (no-op for mock).""" + pass + + async def send_realtime(self, blob: types.Blob): + """Send realtime data (no-op for mock).""" + pass + + async def receive(self) -> AsyncGenerator[LlmResponse, None]: + """Yield each of the pre-defined LlmResponses.""" + for response in self.llm_responses: + yield response + + async def close(self): + """Close connection (no-op for mock).""" + pass diff --git a/tests/unittests/testing_utils.py b/tests/unittests/testing_utils.py index 0bc557e931..12fdcc7fae 100644 --- a/tests/unittests/testing_utils.py +++ b/tests/unittests/testing_utils.py @@ -12,404 +12,44 @@ # See the License for the specific language governing permissions and # limitations under the License. -import asyncio -import contextlib -from typing import AsyncGenerator -from typing import Generator -from typing import Optional -from typing import Union - -from google.adk.agents.invocation_context import InvocationContext -from google.adk.agents.live_request_queue import LiveRequestQueue -from google.adk.agents.llm_agent import Agent -from google.adk.agents.llm_agent import LlmAgent -from google.adk.agents.run_config import RunConfig -from google.adk.apps.app import App -from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService -from google.adk.events.event import Event -from google.adk.memory.in_memory_memory_service import InMemoryMemoryService -from google.adk.models.base_llm import BaseLlm -from google.adk.models.base_llm_connection import BaseLlmConnection -from google.adk.models.llm_request import LlmRequest -from google.adk.models.llm_response import LlmResponse -from google.adk.plugins.base_plugin import BasePlugin -from google.adk.plugins.plugin_manager import PluginManager -from google.adk.runners import InMemoryRunner as AfInMemoryRunner -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 - - -def create_test_agent(name: str = 'test_agent') -> LlmAgent: - """Create a simple test agent for use in unit tests. - - Args: - name: The name of the test agent. - - Returns: - A configured LlmAgent instance suitable for testing. - """ - return LlmAgent(name=name) - - -class UserContent(types.Content): - - def __init__(self, text_or_part: str): - parts = [ - types.Part.from_text(text=text_or_part) - if isinstance(text_or_part, str) - else text_or_part - ] - super().__init__(role='user', parts=parts) - - -class ModelContent(types.Content): - - def __init__(self, parts: list[types.Part]): - super().__init__(role='model', parts=parts) - - -async def create_invocation_context( - agent: Agent, - user_content: str = '', - run_config: RunConfig = None, - plugins: list[BasePlugin] = [], -): - invocation_id = 'test_id' - artifact_service = InMemoryArtifactService() - session_service = InMemorySessionService() - memory_service = InMemoryMemoryService() - invocation_context = InvocationContext( - artifact_service=artifact_service, - session_service=session_service, - memory_service=memory_service, - plugin_manager=PluginManager(plugins=plugins), - invocation_id=invocation_id, - agent=agent, - session=await session_service.create_session( - app_name='test_app', user_id='test_user' - ), - user_content=types.Content( - role='user', parts=[types.Part.from_text(text=user_content)] - ), - run_config=run_config or RunConfig(), - ) - if user_content: - append_user_content( - invocation_context, [types.Part.from_text(text=user_content)] - ) - return invocation_context - - -def append_user_content( - invocation_context: InvocationContext, parts: list[types.Part] -) -> Event: - session = invocation_context.session - event = Event( - invocation_id=invocation_context.invocation_id, - author='user', - content=types.Content(role='user', parts=parts), - ) - session.events.append(event) - return event - - -# Extracts the contents from the events and transform them into a list of -# (author, simplified_content) tuples. -def simplify_events(events: list[Event]) -> list[(str, types.Part)]: - return [ - (event.author, simplify_content(event.content)) - for event in events - if event.content - ] - - -END_OF_AGENT = 'end_of_agent' - - -# Extracts the contents from the events and transform them into a list of -# (author, simplified_content OR AgentState OR "end_of_agent") tuples. -# -# Could be used to compare events for testing resumability. -def simplify_resumable_app_events( - events: list[Event], -) -> list[(str, Union[types.Part, str])]: - results = [] - for event in events: - if event.content: - results.append((event.author, simplify_content(event.content))) - elif event.actions.end_of_agent: - results.append((event.author, END_OF_AGENT)) - elif event.actions.agent_state is not None: - results.append((event.author, event.actions.agent_state)) - return results - - -# Simplifies the contents into a list of (author, simplified_content) tuples. -def simplify_contents(contents: list[types.Content]) -> list[(str, types.Part)]: - return [(content.role, simplify_content(content)) for content in contents] - - -# Simplifies the content so it's easier to assert. -# - If there is only one part, return part -# - If the only part is pure text, return stripped_text -# - If there are multiple parts, return parts -# - remove function_call_id if it exists -def simplify_content( - content: types.Content, -) -> Union[str, types.Part, list[types.Part]]: - for part in content.parts: - if part.function_call and part.function_call.id: - part.function_call.id = None - if part.function_response and part.function_response.id: - part.function_response.id = None - if len(content.parts) == 1: - if content.parts[0].text: - return content.parts[0].text.strip() - else: - return content.parts[0] - return content.parts - - -def get_user_content(message: types.ContentUnion) -> types.Content: - return message if isinstance(message, types.Content) else UserContent(message) - - -class TestInMemoryRunner(AfInMemoryRunner): - """InMemoryRunner that is tailored for tests, features async run method. - - app_name is hardcoded as InMemoryRunner in the parent class. - """ - - 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' - ) - agen = self.run_async( - user_id=session.user_id, - session_id=session.id, - new_message=get_user_content(new_message), - ) - async with Aclosing(agen): - async for event in agen: - yield event - - -class InMemoryRunner: - """InMemoryRunner that is tailored for tests.""" - - def __init__( - self, - root_agent: Optional[Union[Agent, LlmAgent]] = None, - response_modalities: list[str] = None, - plugins: list[BasePlugin] = [], - app: Optional[App] = None, - ): - """Initializes the InMemoryRunner. - - Args: - root_agent: The root agent to run, won't be used if app is provided. - response_modalities: The response modalities of the runner. - plugins: The plugins to use in the runner, won't be used if app is - provided. - app: The app to use in the runner. - """ - if not app: - self.app_name = 'test_app' - self.root_agent = root_agent - self.runner = Runner( - app_name='test_app', - agent=root_agent, - artifact_service=InMemoryArtifactService(), - session_service=InMemorySessionService(), - memory_service=InMemoryMemoryService(), - plugins=plugins, - ) - else: - self.app_name = app.name - self.root_agent = app.root_agent - self.runner = Runner( - app=app, - artifact_service=InMemoryArtifactService(), - session_service=InMemorySessionService(), - memory_service=InMemoryMemoryService(), - ) - self.session_id = None - - @property - def session(self) -> Session: - if not self.session_id: - session = self.runner.session_service.create_session_sync( - app_name=self.app_name, user_id='test_user' - ) - self.session_id = session.id - return session - return self.runner.session_service.get_session_sync( - app_name=self.app_name, user_id='test_user', session_id=self.session_id - ) - - def run(self, new_message: types.ContentUnion) -> list[Event]: - return list( - self.runner.run( - user_id=self.session.user_id, - session_id=self.session.id, - new_message=get_user_content(new_message), - ) - ) - - async def run_async( - self, - new_message: Optional[types.ContentUnion] = None, - invocation_id: Optional[str] = None, - ) -> list[Event]: - events = [] - async for event in self.runner.run_async( - user_id=self.session.user_id, - session_id=self.session.id, - invocation_id=invocation_id, - new_message=get_user_content(new_message) if new_message else None, - ): - events.append(event) - return events - - def run_live( - self, live_request_queue: LiveRequestQueue, run_config: RunConfig = None - ) -> list[Event]: - collected_responses = [] - - async def consume_responses(session: Session): - run_res = self.runner.run_live( - session=session, - live_request_queue=live_request_queue, - run_config=run_config or RunConfig(), - ) - - async for response in run_res: - collected_responses.append(response) - # When we have enough response, we should return - if len(collected_responses) >= 1: - return - - try: - session = self.session - asyncio.run(consume_responses(session)) - except asyncio.TimeoutError: - print('Returning any partial results collected so far.') - - return collected_responses - - -class MockModel(BaseLlm): - model: str = 'mock' - - requests: list[LlmRequest] = [] - responses: list[LlmResponse] - error: Union[Exception, None] = None - response_index: int = -1 - - @classmethod - def create( - cls, - responses: Union[ - list[types.Part], list[LlmResponse], list[str], list[list[types.Part]] - ], - error: Union[Exception, None] = None, - ): - if error and not responses: - return cls(responses=[], error=error) - if not responses: - return cls(responses=[]) - elif isinstance(responses[0], LlmResponse): - # responses is list[LlmResponse] - return cls(responses=responses) - else: - responses = [ - LlmResponse(content=ModelContent(item)) - if isinstance(item, list) and isinstance(item[0], types.Part) - # responses is list[list[Part]] - else LlmResponse( - content=ModelContent( - # responses is list[str] or list[Part] - [Part(text=item) if isinstance(item, str) else item] - ) - ) - for item in responses - if item - ] - - return cls(responses=responses) - - @classmethod - @override - def supported_models(cls) -> list[str]: - return ['mock'] - - def generate_content( - self, llm_request: LlmRequest, stream: bool = False - ) -> Generator[LlmResponse, None, None]: - if self.error is not None: - raise self.error - # Increasement of the index has to happen before the yield. - self.response_index += 1 - self.requests.append(llm_request) - # yield LlmResponse(content=self.responses[self.response_index]) - yield self.responses[self.response_index] - - @override - async def generate_content_async( - self, llm_request: LlmRequest, stream: bool = False - ) -> AsyncGenerator[LlmResponse, None]: - if self.error is not None: - raise self.error - # Increasement of the index has to happen before the yield. - self.response_index += 1 - self.requests.append(llm_request) - yield self.responses[self.response_index] - - @contextlib.asynccontextmanager - async def connect(self, llm_request: LlmRequest) -> BaseLlmConnection: - """Creates a live connection to the LLM.""" - self.requests.append(llm_request) - yield MockLlmConnection(self.responses) - - -class MockLlmConnection(BaseLlmConnection): - - def __init__(self, llm_responses: list[LlmResponse]): - self.llm_responses = llm_responses - - async def send_history(self, history: list[types.Content]): - pass - - async def send_content(self, content: types.Content): - pass - - async def send(self, data): - pass - - async def send_realtime(self, blob: types.Blob): - pass - - async def receive(self) -> AsyncGenerator[LlmResponse, None]: - """Yield each of the pre-defined LlmResponses.""" - for response in self.llm_responses: - yield response - - async def close(self): - pass +"""Backward compatibility shim for tests. + +This module re-exports all testing utilities from the new public location +at google.adk.testing. Tests should continue to import from here for now, +but new code should import from google.adk.testing directly. +""" + +# Re-export everything from the new public testing module +from google.adk.testing import append_user_content +from google.adk.testing import create_invocation_context +from google.adk.testing import create_test_agent +from google.adk.testing import END_OF_AGENT +from google.adk.testing import get_user_content +from google.adk.testing import InMemoryRunner +from google.adk.testing import MockLlmConnection +from google.adk.testing import MockModel +from google.adk.testing import ModelContent +from google.adk.testing import simplify_content +from google.adk.testing import simplify_contents +from google.adk.testing import simplify_events +from google.adk.testing import simplify_resumable_app_events +from google.adk.testing import TestInMemoryRunner +from google.adk.testing import UserContent + +__all__ = [ + 'MockModel', + 'MockLlmConnection', + 'InMemoryRunner', + 'TestInMemoryRunner', + 'create_test_agent', + 'create_invocation_context', + 'append_user_content', + 'simplify_events', + 'simplify_resumable_app_events', + 'simplify_contents', + 'simplify_content', + 'get_user_content', + 'UserContent', + 'ModelContent', + 'END_OF_AGENT', +] From c0759c1ea894201362210cd44709c7693df33456 Mon Sep 17 00:00:00 2001 From: Alexander Bondarenko Date: Fri, 21 Nov 2025 12:08:13 +0100 Subject: [PATCH 2/5] chore: removed unnecessary readme --- src/google/adk/testing/README.md | 245 ------------------------------- 1 file changed, 245 deletions(-) delete mode 100644 src/google/adk/testing/README.md diff --git a/src/google/adk/testing/README.md b/src/google/adk/testing/README.md deleted file mode 100644 index 7456ec2b71..0000000000 --- a/src/google/adk/testing/README.md +++ /dev/null @@ -1,245 +0,0 @@ -# ADK Testing Utilities - -This module provides testing utilities for developers building agents with the Agent Development Kit (ADK). These tools make it easy to write unit tests for your agents without requiring real LLM API calls. - -## Quick Start - -```python -from google.adk import Agent -from google.adk.testing import MockModel, InMemoryRunner - -# Create an agent with a mock model -agent = Agent( - name="test_agent", - model=MockModel.create(responses=["Hello, I'm a test response!"]), - instruction="You are a helpful assistant." -) - -# Run the agent in a test environment -runner = InMemoryRunner(root_agent=agent) -events = runner.run("Hi there!") - -# Assert on the results -assert len(events) > 0 -``` - -## Key Components - -### MockModel - -A mock LLM that returns pre-defined responses instead of calling a real API. This makes tests fast, deterministic, and doesn't require API keys. - -```python -from google.adk.testing import MockModel - -# Simple text responses -mock = MockModel.create(responses=["Response 1", "Response 2"]) - -# Multiple responses for a conversation -mock = MockModel.create(responses=[ - "First response", - "Second response", - "Third response" -]) - -# Mock an error -mock = MockModel.create(responses=[], error=ValueError("API Error")) -``` - -### InMemoryRunner - -A test runner that uses in-memory services for fast, isolated tests. - -```python -from google.adk.testing import InMemoryRunner - -runner = InMemoryRunner(root_agent=agent) - -# Synchronous execution -events = runner.run("Test message") - -# Async execution -events = await runner.run_async("Test message") - -# Access the session -session = runner.session -``` - -### Helper Functions - -#### simplify_events() - -Simplify events for easier assertions in tests: - -```python -from google.adk.testing import simplify_events - -events = runner.run("Hello") -simplified = simplify_events(events) - -# Returns list of (author, simplified_content) tuples -# Easier to assert: [("user", "Hello"), ("test_agent", "Hi there!")] -``` - -#### create_test_agent() - -Create a basic test agent quickly: - -```python -from google.adk.testing import create_test_agent - -agent = create_test_agent(name="my_test") -``` - -#### create_invocation_context() - -Create a test invocation context for testing agent components: - -```python -from google.adk.testing import create_invocation_context - -context = await create_invocation_context( - agent=agent, - user_content="Test message", - run_config=RunConfig(), - plugins=[] -) -``` - -## Complete Example - -```python -import pytest -from google.adk import Agent -from google.adk.testing import MockModel, InMemoryRunner, simplify_events - -@pytest.fixture -def mock_agent(): - """Create a test agent with mock responses.""" - mock_model = MockModel.create(responses=[ - "I can help with that!", - "Here's the information you requested.", - "Is there anything else?" - ]) - - return Agent( - name="test_assistant", - model=mock_model, - instruction="You are a helpful assistant." - ) - -def test_basic_conversation(mock_agent): - """Test a basic conversation flow.""" - runner = InMemoryRunner(root_agent=mock_agent) - - # First message - events = runner.run("Can you help me?") - simplified = simplify_events(events) - - assert len(simplified) >= 2 - assert simplified[0][0] == "user" - assert simplified[0][1] == "Can you help me?" - - # Second message in same session - events = runner.run("Thanks!") - simplified = simplify_events(events) - - assert len(simplified) >= 2 - -@pytest.mark.asyncio -async def test_async_execution(mock_agent): - """Test async agent execution.""" - runner = InMemoryRunner(root_agent=mock_agent) - - events = await runner.run_async("Hello!") - assert len(events) > 0 - -def test_error_handling(): - """Test error handling with mock models.""" - error_model = MockModel.create( - responses=[], - error=ValueError("API unavailable") - ) - - agent = Agent(name="error_agent", model=error_model) - runner = InMemoryRunner(root_agent=agent) - - with pytest.raises(ValueError, match="API unavailable"): - runner.run("This will fail") -``` - -## API Reference - -### MockModel - -- `MockModel.create(responses, error=None)` - Create a mock model with pre-defined responses -- `model.requests` - List of all LlmRequest objects sent to the mock -- `model.responses` - List of LlmResponse objects the mock will return -- `model.response_index` - Current position in the responses list - -### InMemoryRunner - -- `InMemoryRunner(root_agent, plugins=[], app=None)` - Create a test runner -- `runner.run(message)` - Run agent synchronously, returns list of Events -- `runner.run_async(message, invocation_id=None)` - Run agent asynchronously -- `runner.session` - Access the current test session - -### Helper Functions - -- `create_test_agent(name)` - Create a simple test agent -- `create_invocation_context(agent, user_content, run_config, plugins)` - Create test context -- `simplify_events(events)` - Simplify events for assertions -- `simplify_content(content)` - Simplify content for assertions -- `append_user_content(context, parts)` - Add user content to context - -## Best Practices - -1. **Use MockModel for unit tests** - Fast, deterministic, no API calls needed -2. **Use InMemoryRunner** - Isolated test environment with in-memory services -3. **Use simplify_events() for assertions** - Makes test assertions cleaner -4. **Test both sync and async paths** - If your agent supports both -5. **Test error handling** - Use MockModel.create() with error parameter -6. **Keep tests isolated** - Each test should use its own runner instance - -## Integration Testing - -For integration tests with real LLM APIs, use the regular `Runner` instead: - -```python -from google.adk import Agent, Runner -from google.adk.sessions import InMemorySessionService - -# Real agent for integration testing -agent = Agent( - name="real_agent", - model="gemini-2.0-flash-exp", - instruction="You are a helpful assistant." -) - -runner = Runner( - agent=agent, - session_service=InMemorySessionService() -) - -# This will make real API calls -events = list(runner.run( - user_id="test_user", - session_id="test_session", - new_message="Hello!" -)) -``` - -## Migration Guide - -If you were previously using internal testing utilities from `tests/unittests/testing_utils.py`, you can now import from the public module: - -```python -# Old (internal) -from tests.unittests.testing_utils import MockModel, InMemoryRunner - -# New (public) -from google.adk.testing import MockModel, InMemoryRunner -``` - -All functionality remains the same - this just makes the utilities officially supported and available via the PyPI package. - From 5becf3c3aac5912d91c33da9668cbfed33ece0f4 Mon Sep 17 00:00:00 2001 From: Alexander Bondarenko Date: Fri, 21 Nov 2025 12:15:46 +0100 Subject: [PATCH 3/5] chore: removed unnecessary comments --- src/google/adk/testing/testing_utils.py | 214 +----------------------- 1 file changed, 5 insertions(+), 209 deletions(-) diff --git a/src/google/adk/testing/testing_utils.py b/src/google/adk/testing/testing_utils.py index 758bb499c0..158deeb82e 100644 --- a/src/google/adk/testing/testing_utils.py +++ b/src/google/adk/testing/testing_utils.py @@ -12,23 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Public testing utilities for ADK users. - -This module provides utilities for testing agents built with the Agent -Development Kit (ADK). It includes mock models, test runners, and helper -functions to make testing agents easier. - -Example: - >>> from google.adk.testing import MockModel, InMemoryRunner, create_test_agent - >>> agent = create_test_agent(name="my_test_agent") - >>> mock_model = MockModel.create(responses=["Hello, world!"]) - >>> agent.model = mock_model - >>> runner = InMemoryRunner(root_agent=agent) - >>> events = runner.run("Hi there!") -""" - -from __future__ import annotations - import asyncio import contextlib from typing import AsyncGenerator @@ -74,7 +57,6 @@ def create_test_agent(name: str = 'test_agent') -> LlmAgent: class UserContent(types.Content): - """Helper to create user content for testing.""" def __init__(self, text_or_part: str): parts = [ @@ -86,7 +68,6 @@ def __init__(self, text_or_part: str): class ModelContent(types.Content): - """Helper to create model content for testing.""" def __init__(self, parts: list[types.Part]): super().__init__(role='model', parts=parts) @@ -98,17 +79,6 @@ async def create_invocation_context( run_config: RunConfig = None, plugins: list[BasePlugin] = [], ): - """Create an invocation context for testing agent components. - - Args: - agent: The agent to create context for. - user_content: Optional user message content. - run_config: Optional run configuration. - plugins: Optional list of plugins to use. - - Returns: - An InvocationContext configured for testing. - """ invocation_id = 'test_id' artifact_service = InMemoryArtifactService() session_service = InMemorySessionService() @@ -138,15 +108,6 @@ async def create_invocation_context( def append_user_content( invocation_context: InvocationContext, parts: list[types.Part] ) -> Event: - """Append user content to an invocation context's session. - - Args: - invocation_context: The context to append to. - parts: Content parts to append. - - Returns: - The created Event. - """ session = invocation_context.session event = Event( invocation_id=invocation_context.invocation_id, @@ -160,14 +121,6 @@ def append_user_content( # Extracts the contents from the events and transform them into a list of # (author, simplified_content) tuples. def simplify_events(events: list[Event]) -> list[(str, types.Part)]: - """Simplify events for easier assertion in tests. - - Args: - events: List of events to simplify. - - Returns: - List of (author, simplified_content) tuples. - """ return [ (event.author, simplify_content(event.content)) for event in events @@ -185,14 +138,6 @@ def simplify_events(events: list[Event]) -> list[(str, types.Part)]: def simplify_resumable_app_events( events: list[Event], ) -> list[(str, Union[types.Part, str])]: - """Simplify events including agent state for resumability testing. - - Args: - events: List of events to simplify. - - Returns: - List of (author, simplified_content/state) tuples. - """ results = [] for event in events: if event.content: @@ -206,14 +151,6 @@ def simplify_resumable_app_events( # Simplifies the contents into a list of (author, simplified_content) tuples. def simplify_contents(contents: list[types.Content]) -> list[(str, types.Part)]: - """Simplify contents for easier assertion in tests. - - Args: - contents: List of contents to simplify. - - Returns: - List of (role, simplified_content) tuples. - """ return [(content.role, simplify_content(content)) for content in contents] @@ -225,19 +162,6 @@ def simplify_contents(contents: list[types.Content]) -> list[(str, types.Part)]: def simplify_content( content: types.Content, ) -> Union[str, types.Part, list[types.Part]]: - """Simplify content for easier assertion in tests. - - Removes function call IDs and returns simplified representation: - - Single text part: returns stripped text string - - Single non-text part: returns the part - - Multiple parts: returns list of parts - - Args: - content: Content to simplify. - - Returns: - Simplified content representation. - """ for part in content.parts: if part.function_call and part.function_call.id: part.function_call.id = None @@ -252,35 +176,19 @@ def simplify_content( def get_user_content(message: types.ContentUnion) -> types.Content: - """Convert a message to user Content. - - Args: - message: Either a Content object or string message. - - Returns: - A Content object with user role. - """ return message if isinstance(message, types.Content) else UserContent(message) class TestInMemoryRunner(AfInMemoryRunner): - """InMemoryRunner tailored for async tests. + """InMemoryRunner that is tailored for tests, features async run method. - This runner extends the base InMemoryRunner with async run methods - suitable for unit testing. + app_name is hardcoded as InMemoryRunner in the parent class. """ async def run_async_with_new_session( self, new_message: types.ContentUnion ) -> list[Event]: - """Run agent with a new session and collect all events. - - Args: - new_message: The user message to send. - Returns: - List of all events generated during execution. - """ collected_events: list[Event] = [] async for event in self.run_async_with_new_session_agen(new_message): collected_events.append(event) @@ -290,14 +198,6 @@ async def run_async_with_new_session( async def run_async_with_new_session_agen( self, new_message: types.ContentUnion ) -> AsyncGenerator[Event, None]: - """Run agent with a new session as an async generator. - - Args: - new_message: The user message to send. - - Yields: - Events as they are generated. - """ session = await self.session_service.create_session( app_name='InMemoryRunner', user_id='test_user' ) @@ -312,17 +212,7 @@ async def run_async_with_new_session_agen( class InMemoryRunner: - """Test runner with in-memory services for testing agents. - - This runner is designed specifically for testing, providing easy access - to session state and in-memory services. - - Example: - >>> agent = create_test_agent() - >>> runner = InMemoryRunner(root_agent=agent) - >>> events = runner.run("Hello!") - >>> assert len(events) > 0 - """ + """InMemoryRunner that is tailored for tests.""" def __init__( self, @@ -364,7 +254,6 @@ def __init__( @property def session(self) -> Session: - """Get or create the test session.""" if not self.session_id: session = self.runner.session_service.create_session_sync( app_name=self.app_name, user_id='test_user' @@ -376,14 +265,6 @@ def session(self) -> Session: ) def run(self, new_message: types.ContentUnion) -> list[Event]: - """Run the agent synchronously and return all events. - - Args: - new_message: The user message to send. - - Returns: - List of all events generated during execution. - """ return list( self.runner.run( user_id=self.session.user_id, @@ -397,15 +278,6 @@ async def run_async( new_message: Optional[types.ContentUnion] = None, invocation_id: Optional[str] = None, ) -> list[Event]: - """Run the agent asynchronously and return all events. - - Args: - new_message: The user message to send. - invocation_id: Optional invocation ID for tracking. - - Returns: - List of all events generated during execution. - """ events = [] async for event in self.runner.run_async( user_id=self.session.user_id, @@ -419,15 +291,6 @@ async def run_async( def run_live( self, live_request_queue: LiveRequestQueue, run_config: RunConfig = None ) -> list[Event]: - """Run the agent in live mode (bi-directional streaming). - - Args: - live_request_queue: Queue for live requests. - run_config: Optional run configuration. - - Returns: - List of events from the live session. - """ collected_responses = [] async def consume_responses(session: Session): @@ -453,17 +316,6 @@ async def consume_responses(session: Session): class MockModel(BaseLlm): - """Mock LLM model for testing without real API calls. - - This model returns pre-defined responses instead of calling a real LLM API, - making tests fast, deterministic, and not requiring API keys. - - Example: - >>> mock = MockModel.create(responses=["Response 1", "Response 2"]) - >>> agent = LlmAgent(name="test", model=mock) - >>> # Agent will return "Response 1" on first call, "Response 2" on second - """ - model: str = 'mock' requests: list[LlmRequest] = [] @@ -479,19 +331,6 @@ def create( ], error: Union[Exception, None] = None, ): - """Create a MockModel with pre-defined responses. - - Args: - responses: List of responses to return. Can be: - - list[str]: Simple text responses - - list[Part]: Content parts - - list[list[Part]]: Multi-part responses - - list[LlmResponse]: Full response objects - error: Optional exception to raise instead of returning responses. - - Returns: - A configured MockModel instance. - """ if error and not responses: return cls(responses=[], error=error) if not responses: @@ -519,24 +358,11 @@ def create( @classmethod @override def supported_models(cls) -> list[str]: - """Return list of supported model names.""" return ['mock'] def generate_content( self, llm_request: LlmRequest, stream: bool = False ) -> Generator[LlmResponse, None, None]: - """Generate content synchronously. - - Args: - llm_request: The request to process. - stream: Whether to stream the response (ignored for mock). - - Yields: - The next pre-defined response. - - Raises: - Exception: If an error was configured. - """ if self.error is not None: raise self.error # Increasement of the index has to happen before the yield. @@ -549,18 +375,6 @@ def generate_content( async def generate_content_async( self, llm_request: LlmRequest, stream: bool = False ) -> AsyncGenerator[LlmResponse, None]: - """Generate content asynchronously. - - Args: - llm_request: The request to process. - stream: Whether to stream the response (ignored for mock). - - Yields: - The next pre-defined response. - - Raises: - Exception: If an error was configured. - """ if self.error is not None: raise self.error # Increasement of the index has to happen before the yield. @@ -570,43 +384,26 @@ async def generate_content_async( @contextlib.asynccontextmanager async def connect(self, llm_request: LlmRequest) -> BaseLlmConnection: - """Create a live connection to the mock LLM. - - Args: - llm_request: The request to process. - - Yields: - A mock connection that returns pre-defined responses. - """ + """Creates a live connection to the LLM.""" self.requests.append(llm_request) yield MockLlmConnection(self.responses) class MockLlmConnection(BaseLlmConnection): - """Mock LLM connection for live/streaming tests.""" def __init__(self, llm_responses: list[LlmResponse]): - """Initialize with pre-defined responses. - - Args: - llm_responses: Responses to return. - """ self.llm_responses = llm_responses async def send_history(self, history: list[types.Content]): - """Send history (no-op for mock).""" pass async def send_content(self, content: types.Content): - """Send content (no-op for mock).""" pass async def send(self, data): - """Send data (no-op for mock).""" pass async def send_realtime(self, blob: types.Blob): - """Send realtime data (no-op for mock).""" pass async def receive(self) -> AsyncGenerator[LlmResponse, None]: @@ -615,5 +412,4 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]: yield response async def close(self): - """Close connection (no-op for mock).""" - pass + pass \ No newline at end of file From f6446e6095b8329aa8fd48e4cb376b6d4784f1b8 Mon Sep 17 00:00:00 2001 From: Alexander Bondarenko Date: Fri, 21 Nov 2025 12:26:00 +0100 Subject: [PATCH 4/5] fix(testing): Export commonly used types for backward compatibility - Export RunConfig, Event, LlmRequest, LlmResponse, Session from testing module - Tests access these types via testing_utils.Type (e.g., testing_utils.RunConfig) - Use relative imports in src/ per ADK style guide - Maintains full backward compatibility for all existing tests --- src/google/adk/testing/__init__.py | 23 +++++++++++++++++------ tests/unittests/testing_utils.py | 13 +++++++++++++ 2 files changed, 30 insertions(+), 6 deletions(-) diff --git a/src/google/adk/testing/__init__.py b/src/google/adk/testing/__init__.py index 6ceca0c46a..3703211723 100644 --- a/src/google/adk/testing/__init__.py +++ b/src/google/adk/testing/__init__.py @@ -55,21 +55,26 @@ from __future__ import annotations -from .testing_utils import append_user_content -from .testing_utils import create_invocation_context -from .testing_utils import create_test_agent +from ..agents.run_config import RunConfig +from ..events.event import Event +from ..models.llm_request import LlmRequest +from ..models.llm_response import LlmResponse +from ..sessions.session import Session from .testing_utils import END_OF_AGENT -from .testing_utils import get_user_content from .testing_utils import InMemoryRunner from .testing_utils import MockLlmConnection from .testing_utils import MockModel from .testing_utils import ModelContent +from .testing_utils import TestInMemoryRunner +from .testing_utils import UserContent +from .testing_utils import append_user_content +from .testing_utils import create_invocation_context +from .testing_utils import create_test_agent +from .testing_utils import get_user_content from .testing_utils import simplify_content from .testing_utils import simplify_contents from .testing_utils import simplify_events from .testing_utils import simplify_resumable_app_events -from .testing_utils import TestInMemoryRunner -from .testing_utils import UserContent __all__ = [ 'MockModel', @@ -87,4 +92,10 @@ 'UserContent', 'ModelContent', 'END_OF_AGENT', + # Commonly used types + 'RunConfig', + 'Event', + 'LlmRequest', + 'LlmResponse', + 'Session', ] diff --git a/tests/unittests/testing_utils.py b/tests/unittests/testing_utils.py index 12fdcc7fae..fd3e6520f8 100644 --- a/tests/unittests/testing_utils.py +++ b/tests/unittests/testing_utils.py @@ -36,6 +36,13 @@ from google.adk.testing import TestInMemoryRunner from google.adk.testing import UserContent +# Re-export commonly used types that tests access via testing_utils +from google.adk.agents.run_config import RunConfig +from google.adk.events.event import Event +from google.adk.models.llm_request import LlmRequest +from google.adk.models.llm_response import LlmResponse +from google.adk.sessions.session import Session + __all__ = [ 'MockModel', 'MockLlmConnection', @@ -52,4 +59,10 @@ 'UserContent', 'ModelContent', 'END_OF_AGENT', + # Commonly used types + 'RunConfig', + 'Event', + 'LlmRequest', + 'LlmResponse', + 'Session', ] From f21f5cdf27e72e048c22107e7408f4a53ce9590b Mon Sep 17 00:00:00 2001 From: Alexander Bondarenko Date: Wed, 3 Dec 2025 16:06:55 +0100 Subject: [PATCH 5/5] fix: add missing __future__ annotations import and fix formatting - Add 'from __future__ import annotations' to testing_utils.py - Fix import sorting in testing files and samples - Fix pyink formatting (add missing newline at EOF) --- contributing/samples/gepa/experiment.py | 1 - contributing/samples/gepa/run_experiment.py | 1 - src/google/adk/testing/__init__.py | 12 ++++++------ src/google/adk/testing/testing_utils.py | 4 +++- tests/unittests/testing_utils.py | 13 ++++++------- 5 files changed, 15 insertions(+), 16 deletions(-) diff --git a/contributing/samples/gepa/experiment.py b/contributing/samples/gepa/experiment.py index 2f5d03a772..f68b349d9c 100644 --- a/contributing/samples/gepa/experiment.py +++ b/contributing/samples/gepa/experiment.py @@ -43,7 +43,6 @@ from tau_bench.types import EnvRunResult from tau_bench.types import RunConfig import tau_bench_agent as tau_bench_agent_lib - import utils diff --git a/contributing/samples/gepa/run_experiment.py b/contributing/samples/gepa/run_experiment.py index cfd850b3a3..1bc4ee58c8 100644 --- a/contributing/samples/gepa/run_experiment.py +++ b/contributing/samples/gepa/run_experiment.py @@ -25,7 +25,6 @@ from absl import flags import experiment from google.genai import types - import utils _OUTPUT_DIR = flags.DEFINE_string( diff --git a/src/google/adk/testing/__init__.py b/src/google/adk/testing/__init__.py index 3703211723..360f7efc34 100644 --- a/src/google/adk/testing/__init__.py +++ b/src/google/adk/testing/__init__.py @@ -60,21 +60,21 @@ from ..models.llm_request import LlmRequest from ..models.llm_response import LlmResponse from ..sessions.session import Session +from .testing_utils import append_user_content +from .testing_utils import create_invocation_context +from .testing_utils import create_test_agent from .testing_utils import END_OF_AGENT +from .testing_utils import get_user_content from .testing_utils import InMemoryRunner from .testing_utils import MockLlmConnection from .testing_utils import MockModel from .testing_utils import ModelContent -from .testing_utils import TestInMemoryRunner -from .testing_utils import UserContent -from .testing_utils import append_user_content -from .testing_utils import create_invocation_context -from .testing_utils import create_test_agent -from .testing_utils import get_user_content from .testing_utils import simplify_content from .testing_utils import simplify_contents from .testing_utils import simplify_events from .testing_utils import simplify_resumable_app_events +from .testing_utils import TestInMemoryRunner +from .testing_utils import UserContent __all__ = [ 'MockModel', diff --git a/src/google/adk/testing/testing_utils.py b/src/google/adk/testing/testing_utils.py index 158deeb82e..0cc00c7272 100644 --- a/src/google/adk/testing/testing_utils.py +++ b/src/google/adk/testing/testing_utils.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import asyncio import contextlib from typing import AsyncGenerator @@ -412,4 +414,4 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]: yield response async def close(self): - pass \ No newline at end of file + pass diff --git a/tests/unittests/testing_utils.py b/tests/unittests/testing_utils.py index fd3e6520f8..5d2a6fbaa3 100644 --- a/tests/unittests/testing_utils.py +++ b/tests/unittests/testing_utils.py @@ -19,6 +19,12 @@ but new code should import from google.adk.testing directly. """ +# Re-export commonly used types that tests access via testing_utils +from google.adk.agents.run_config import RunConfig +from google.adk.events.event import Event +from google.adk.models.llm_request import LlmRequest +from google.adk.models.llm_response import LlmResponse +from google.adk.sessions.session import Session # Re-export everything from the new public testing module from google.adk.testing import append_user_content from google.adk.testing import create_invocation_context @@ -36,13 +42,6 @@ from google.adk.testing import TestInMemoryRunner from google.adk.testing import UserContent -# Re-export commonly used types that tests access via testing_utils -from google.adk.agents.run_config import RunConfig -from google.adk.events.event import Event -from google.adk.models.llm_request import LlmRequest -from google.adk.models.llm_response import LlmResponse -from google.adk.sessions.session import Session - __all__ = [ 'MockModel', 'MockLlmConnection',