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
3271 lines
111 KiB
Python
3271 lines
111 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.
|
|
|
|
import copy
|
|
import json
|
|
from pathlib import Path
|
|
import tempfile
|
|
from unittest.mock import AsyncMock
|
|
from unittest.mock import create_autospec
|
|
from unittest.mock import Mock
|
|
from unittest.mock import patch
|
|
|
|
from a2a.client import Client as A2AClient
|
|
from a2a.client.client import ClientConfig
|
|
from a2a.client.client_factory import ClientFactory
|
|
from a2a.types import AgentCapabilities
|
|
from a2a.types import AgentCard
|
|
from a2a.types import AgentSkill
|
|
from a2a.types import Artifact
|
|
from a2a.types import Message as A2AMessage
|
|
from a2a.types import Task as A2ATask
|
|
from a2a.types import TaskArtifactUpdateEvent
|
|
from a2a.types import TaskStatus as A2ATaskStatus
|
|
from a2a.types import TaskStatusUpdateEvent
|
|
from google.adk.a2a import _compat
|
|
from google.adk.a2a.agent import ParametersConfig
|
|
from google.adk.a2a.agent import RequestInterceptor
|
|
from google.adk.a2a.agent.config import A2aRemoteAgentConfig
|
|
from google.adk.a2a.agent.utils import execute_after_request_interceptors
|
|
from google.adk.a2a.agent.utils import execute_before_request_interceptors
|
|
from google.adk.agents.invocation_context import InvocationContext
|
|
from google.adk.agents.remote_a2a_agent import A2A_METADATA_PREFIX
|
|
from google.adk.agents.remote_a2a_agent import AgentCardResolutionError
|
|
from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
|
|
import google.adk.agents.remote_a2a_agent as remote_a2a_agent
|
|
from google.adk.events.event import Event
|
|
from google.adk.sessions.session import Session
|
|
from google.genai import types as genai_types
|
|
import httpx
|
|
import pytest
|
|
|
|
|
|
def _make_agent_card(
|
|
name="test-agent",
|
|
url="https://example.com/rpc",
|
|
description="Test agent",
|
|
*,
|
|
version="1.0",
|
|
skills=None,
|
|
**kwargs,
|
|
):
|
|
"""Build an AgentCard version-agnostically for tests."""
|
|
|
|
if skills is None:
|
|
skills = []
|
|
if _compat.IS_A2A_V1:
|
|
card = _compat.parse_agent_card({
|
|
"name": name,
|
|
"description": description,
|
|
"version": version,
|
|
"supported_interfaces": [{"url": url, "protocol_binding": "JSONRPC"}],
|
|
"default_input_modes": ["text/plain"],
|
|
"default_output_modes": ["text/plain"],
|
|
})
|
|
for skill in skills:
|
|
card.skills.append(skill)
|
|
return card
|
|
else:
|
|
return AgentCard(
|
|
name=name,
|
|
url=url,
|
|
description=description,
|
|
version=version,
|
|
capabilities=AgentCapabilities(),
|
|
skills=skills,
|
|
default_input_modes=["text/plain"],
|
|
default_output_modes=["text/plain"],
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
def _make_stream_message(message: A2AMessage):
|
|
"""Wrap a Message in the shape ``send_message`` yields for the active SDK.
|
|
|
|
On 1.x ``send_message`` yields ``StreamResponse`` proto objects; on 0.3.x it
|
|
yields the bare ``Message``. ``_compat.make_stream_normalizer`` collapses both
|
|
back to the legacy shape, so tests build the version-correct raw item here.
|
|
"""
|
|
if _compat.IS_A2A_V1:
|
|
from a2a.types import StreamResponse
|
|
|
|
resp = StreamResponse()
|
|
resp.message.CopyFrom(message)
|
|
return resp
|
|
return message
|
|
|
|
|
|
# Helper function to create a proper AgentCard for testing
|
|
def create_test_agent_card(
|
|
name: str = "test-agent",
|
|
url: str = "https://example.com/rpc",
|
|
description: str = "Test agent",
|
|
) -> AgentCard:
|
|
"""Create a test AgentCard with all required fields."""
|
|
return _make_agent_card(
|
|
name=name,
|
|
url=url,
|
|
description=description,
|
|
version="1.0",
|
|
skills=[
|
|
AgentSkill(
|
|
id="test-skill",
|
|
name="Test Skill",
|
|
description="A test skill",
|
|
tags=["test"],
|
|
)
|
|
],
|
|
)
|
|
|
|
|
|
class TestRemoteA2aAgentInit:
|
|
"""Test RemoteA2aAgent initialization and validation."""
|
|
|
|
def test_init_with_agent_card_object(self):
|
|
"""Test initialization with AgentCard object."""
|
|
agent_card = create_test_agent_card()
|
|
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent", agent_card=agent_card, description="Test description"
|
|
)
|
|
|
|
assert agent.name == "test_agent"
|
|
assert agent.description == "Test description"
|
|
assert agent._agent_card == agent_card
|
|
assert agent._agent_card_source is None
|
|
assert agent._httpx_client_needs_cleanup is True
|
|
assert agent._is_resolved is False
|
|
|
|
def test_init_with_url_string(self):
|
|
"""Test initialization with URL string."""
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent", agent_card="https://example.com/agent.json"
|
|
)
|
|
|
|
assert agent.name == "test_agent"
|
|
assert agent._agent_card is None
|
|
assert agent._agent_card_source == "https://example.com/agent.json"
|
|
|
|
def test_init_with_file_path(self):
|
|
"""Test initialization with file path."""
|
|
agent = RemoteA2aAgent(name="test_agent", agent_card="/path/to/agent.json")
|
|
|
|
assert agent.name == "test_agent"
|
|
assert agent._agent_card is None
|
|
assert agent._agent_card_source == "/path/to/agent.json"
|
|
|
|
def test_init_with_shared_httpx_client(self):
|
|
"""Test initialization with shared httpx client."""
|
|
httpx_client = httpx.AsyncClient()
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card="https://example.com/agent.json",
|
|
httpx_client=httpx_client,
|
|
)
|
|
|
|
assert agent._httpx_client is not None
|
|
assert agent._httpx_client_needs_cleanup is False
|
|
|
|
def test_init_with_factory(self):
|
|
"""Test initialization with shared httpx client."""
|
|
httpx_client = httpx.AsyncClient()
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card="https://example.com/agent.json",
|
|
httpx_client=httpx_client,
|
|
)
|
|
|
|
assert agent._httpx_client == httpx_client
|
|
assert agent._httpx_client_needs_cleanup is False
|
|
|
|
def test_init_with_none_agent_card(self):
|
|
"""Test initialization with None agent card raises ValueError."""
|
|
with pytest.raises(ValueError, match="agent_card cannot be None"):
|
|
RemoteA2aAgent(name="test_agent", agent_card=None)
|
|
|
|
def test_init_with_empty_string_agent_card(self):
|
|
"""Test initialization with empty string agent card raises ValueError."""
|
|
with pytest.raises(ValueError, match="agent_card string cannot be empty"):
|
|
RemoteA2aAgent(name="test_agent", agent_card=" ")
|
|
|
|
def test_init_with_invalid_type_agent_card(self):
|
|
"""Test initialization with invalid type agent card raises TypeError."""
|
|
with pytest.raises(TypeError, match="agent_card must be AgentCard"):
|
|
RemoteA2aAgent(name="test_agent", agent_card=123)
|
|
|
|
def test_init_with_custom_timeout(self):
|
|
"""Test initialization with custom timeout."""
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card="https://example.com/agent.json",
|
|
timeout=300.0,
|
|
)
|
|
|
|
assert agent._timeout == 300.0
|
|
|
|
|
|
class TestRemoteA2aAgentResolution:
|
|
"""Test agent card resolution functionality."""
|
|
|
|
def setup_method(self):
|
|
"""Setup test fixtures."""
|
|
self.agent_card_data = {
|
|
"name": "test-agent",
|
|
"url": "https://example.com/rpc",
|
|
"description": "Test agent",
|
|
"version": "1.0",
|
|
"capabilities": {},
|
|
"defaultInputModes": ["text/plain"],
|
|
"defaultOutputModes": ["application/json"],
|
|
"skills": [{
|
|
"id": "test-skill",
|
|
"name": "Test Skill",
|
|
"description": "A test skill",
|
|
"tags": ["test"],
|
|
}],
|
|
}
|
|
self.agent_card = create_test_agent_card()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ensure_httpx_client_creates_new_client(self):
|
|
"""Test that _ensure_httpx_client creates new client when none exists."""
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent", agent_card=create_test_agent_card()
|
|
)
|
|
|
|
client = await agent._ensure_httpx_client()
|
|
|
|
assert client is not None
|
|
assert agent._httpx_client == client
|
|
assert agent._httpx_client_needs_cleanup is True
|
|
|
|
if not _compat.IS_A2A_V1:
|
|
assert agent._a2a_client_factory._config.supported_transports == [
|
|
_compat.TransportProtocol.jsonrpc,
|
|
_compat.TransportProtocol.http_json,
|
|
]
|
|
# 1.x uses supported_protocol_bindings instead.
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ensure_httpx_client_reuses_existing_client(self):
|
|
"""Test that _ensure_httpx_client reuses existing client."""
|
|
existing_client = httpx.AsyncClient()
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card=create_test_agent_card(),
|
|
httpx_client=existing_client,
|
|
)
|
|
|
|
client = await agent._ensure_httpx_client()
|
|
|
|
assert client == existing_client
|
|
assert agent._httpx_client_needs_cleanup is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ensure_factory_reuses_existing_client(self):
|
|
"""Test that _ensure_httpx_client reuses existing client."""
|
|
existing_client = httpx.AsyncClient()
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card=create_test_agent_card(),
|
|
a2a_client_factory=ClientFactory(
|
|
ClientConfig(httpx_client=existing_client),
|
|
),
|
|
)
|
|
|
|
client = await agent._ensure_httpx_client()
|
|
|
|
assert client == existing_client
|
|
assert agent._httpx_client_needs_cleanup is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ensure_httpx_client_updates_factory_with_new_client(self):
|
|
"""Test that _ensure_httpx_client updates factory with new client."""
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card=create_test_agent_card(),
|
|
a2a_client_factory=ClientFactory(
|
|
ClientConfig(httpx_client=None),
|
|
),
|
|
)
|
|
assert agent._a2a_client_factory._config.httpx_client is None
|
|
|
|
client = await agent._ensure_httpx_client()
|
|
|
|
assert client is not None
|
|
assert agent._httpx_client == client
|
|
assert agent._httpx_client_needs_cleanup is True
|
|
assert agent._a2a_client_factory._config.httpx_client == client
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ensure_httpx_client_reregisters_transports_with_new_client(
|
|
self,
|
|
):
|
|
"""Test that _ensure_httpx_client registers transports with new client."""
|
|
factory = ClientFactory(
|
|
ClientConfig(httpx_client=None),
|
|
)
|
|
factory.register("transport_label", lambda: "test")
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card=create_test_agent_card(),
|
|
a2a_client_factory=factory,
|
|
)
|
|
assert agent._a2a_client_factory._config.httpx_client is None
|
|
assert "transport_label" in agent._a2a_client_factory._registry
|
|
|
|
client = await agent._ensure_httpx_client()
|
|
|
|
assert client is not None
|
|
assert agent._httpx_client == client
|
|
assert agent._httpx_client_needs_cleanup is True
|
|
assert agent._a2a_client_factory._config.httpx_client == client
|
|
if not _compat.IS_A2A_V1:
|
|
# On 0.3.x the factory is reconstructed preserving custom
|
|
# transports. On 1.x the factory is recreated fresh with only the
|
|
# standard protocol bindings, so custom transports are not
|
|
# preserved (intended production behavior).
|
|
assert "transport_label" in agent._a2a_client_factory._registry
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_agent_card_from_url_success(self):
|
|
"""Test successful agent card resolution from URL."""
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent", agent_card="https://example.com/agent.json"
|
|
)
|
|
|
|
with patch.object(agent, "_ensure_httpx_client") as mock_ensure_client:
|
|
mock_client = AsyncMock()
|
|
mock_ensure_client.return_value = mock_client
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.A2ACardResolver"
|
|
) as mock_resolver_class:
|
|
mock_resolver = AsyncMock()
|
|
mock_resolver.get_agent_card.return_value = self.agent_card
|
|
mock_resolver_class.return_value = mock_resolver
|
|
|
|
result = await agent._resolve_agent_card_from_url(
|
|
"https://example.com/agent.json"
|
|
)
|
|
|
|
assert result == self.agent_card
|
|
mock_resolver_class.assert_called_once_with(
|
|
httpx_client=mock_client, base_url="https://example.com"
|
|
)
|
|
mock_resolver.get_agent_card.assert_called_once_with(
|
|
relative_card_path="/agent.json"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_agent_card_from_url_invalid_url(self):
|
|
"""Test agent card resolution from invalid URL raises error."""
|
|
agent = RemoteA2aAgent(name="test_agent", agent_card="invalid-url")
|
|
|
|
with pytest.raises(AgentCardResolutionError, match="Invalid URL format"):
|
|
await agent._resolve_agent_card_from_url("invalid-url")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_agent_card_from_file_success(self):
|
|
"""Test successful agent card resolution from file."""
|
|
agent = RemoteA2aAgent(name="test_agent", agent_card="/path/to/agent.json")
|
|
|
|
with tempfile.NamedTemporaryFile(
|
|
mode="w", suffix=".json", delete=False
|
|
) as f:
|
|
json.dump(self.agent_card_data, f)
|
|
temp_path = f.name
|
|
|
|
try:
|
|
result = await agent._resolve_agent_card_from_file(temp_path)
|
|
assert result.name == self.agent_card.name
|
|
assert _compat.agent_card_url(result) == _compat.agent_card_url(
|
|
self.agent_card
|
|
)
|
|
finally:
|
|
Path(temp_path).unlink()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_agent_card_from_file_not_found(self):
|
|
"""Test agent card resolution from nonexistent file raises error."""
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent", agent_card="/path/to/nonexistent.json"
|
|
)
|
|
|
|
with pytest.raises(
|
|
AgentCardResolutionError, match="Agent card file not found"
|
|
):
|
|
await agent._resolve_agent_card_from_file("/path/to/nonexistent.json")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_agent_card_from_file_invalid_json(self):
|
|
"""Test agent card resolution from file with invalid JSON raises error."""
|
|
agent = RemoteA2aAgent(name="test_agent", agent_card="/path/to/agent.json")
|
|
|
|
with tempfile.NamedTemporaryFile(
|
|
mode="w", suffix=".json", delete=False
|
|
) as f:
|
|
f.write("invalid json")
|
|
temp_path = f.name
|
|
|
|
try:
|
|
with pytest.raises(AgentCardResolutionError, match="Invalid JSON"):
|
|
await agent._resolve_agent_card_from_file(temp_path)
|
|
finally:
|
|
Path(temp_path).unlink()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_validate_agent_card_success(self):
|
|
"""Test successful agent card validation."""
|
|
agent_card = create_test_agent_card()
|
|
agent = RemoteA2aAgent(name="test_agent", agent_card=agent_card)
|
|
|
|
# Should not raise any exception
|
|
await agent._validate_agent_card(agent_card)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_validate_agent_card_no_url(self):
|
|
"""Test agent card validation fails when no URL."""
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent", agent_card=create_test_agent_card()
|
|
)
|
|
|
|
invalid_card = _make_agent_card(
|
|
name="test",
|
|
url="", # Empty URL to trigger validation error
|
|
)
|
|
|
|
with pytest.raises(
|
|
AgentCardResolutionError, match="Agent card must have a valid URL"
|
|
):
|
|
await agent._validate_agent_card(invalid_card)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_validate_agent_card_invalid_url(self):
|
|
"""Test agent card validation fails with invalid URL."""
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent", agent_card=create_test_agent_card()
|
|
)
|
|
|
|
invalid_card = _make_agent_card(
|
|
name="test",
|
|
url="invalid-url", # Invalid URL to trigger validation error
|
|
)
|
|
|
|
with pytest.raises(AgentCardResolutionError, match="Invalid RPC URL"):
|
|
await agent._validate_agent_card(invalid_card)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ensure_resolved_with_direct_agent_card(self):
|
|
"""Test _ensure_resolved with direct agent card."""
|
|
agent_card = create_test_agent_card()
|
|
agent = RemoteA2aAgent(name="test_agent", agent_card=agent_card)
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_class:
|
|
mock_client = AsyncMock()
|
|
mock_client_class.return_value = mock_client
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.A2AClientFactory"
|
|
) as mock_factory_class:
|
|
mock_factory = Mock()
|
|
mock_a2a_client = Mock()
|
|
mock_factory.create.return_value = mock_a2a_client
|
|
mock_factory_class.return_value = mock_factory
|
|
|
|
await agent._ensure_resolved()
|
|
|
|
assert agent._is_resolved is True
|
|
assert agent._a2a_client == mock_a2a_client
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ensure_resolved_with_direct_agent_card_with_factory(self):
|
|
"""Test _ensure_resolved with direct agent card."""
|
|
agent_card = create_test_agent_card()
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card=agent_card,
|
|
a2a_client_factory=ClientFactory(
|
|
ClientConfig(),
|
|
),
|
|
)
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_class:
|
|
mock_client = AsyncMock()
|
|
mock_client_class.return_value = mock_client
|
|
|
|
# Rebinding reconstructs the factory through
|
|
# ``_compat.A2AClientFactory``, so the patch must target it there.
|
|
with patch(
|
|
"google.adk.a2a._compat.A2AClientFactory"
|
|
) as mock_factory_class:
|
|
mock_a2a_client = Mock()
|
|
mock_factory = Mock()
|
|
mock_factory.create.return_value = mock_a2a_client
|
|
mock_factory_class.return_value = mock_factory
|
|
|
|
await agent._ensure_resolved()
|
|
|
|
assert agent._is_resolved is True
|
|
assert agent._a2a_client == mock_a2a_client
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ensure_resolved_with_url_source(self):
|
|
"""Test _ensure_resolved with URL source."""
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent", agent_card="https://example.com/agent.json"
|
|
)
|
|
|
|
agent_card = create_test_agent_card()
|
|
with patch.object(agent, "_resolve_agent_card") as mock_resolve:
|
|
mock_resolve.return_value = agent_card
|
|
|
|
with patch.object(agent, "_ensure_httpx_client") as mock_ensure_client:
|
|
mock_client = AsyncMock()
|
|
mock_ensure_client.return_value = mock_client
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.A2AClient"
|
|
) as mock_client_class:
|
|
mock_a2a_client = AsyncMock()
|
|
mock_client_class.return_value = mock_a2a_client
|
|
|
|
await agent._ensure_resolved()
|
|
|
|
assert agent._is_resolved is True
|
|
assert agent._agent_card == agent_card
|
|
assert agent.description == agent_card.description
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ensure_resolved_already_resolved(self):
|
|
"""Test _ensure_resolved when already resolved."""
|
|
agent_card = create_test_agent_card()
|
|
agent = RemoteA2aAgent(name="test_agent", agent_card=agent_card)
|
|
|
|
# Set up as already resolved
|
|
agent._is_resolved = True
|
|
agent._a2a_client = AsyncMock()
|
|
|
|
with patch.object(agent, "_resolve_agent_card") as mock_resolve:
|
|
await agent._ensure_resolved()
|
|
|
|
# Should not call resolution again
|
|
mock_resolve.assert_not_called()
|
|
|
|
|
|
class TestRemoteA2aAgentMessageHandling:
|
|
"""Test message handling functionality."""
|
|
|
|
def setup_method(self):
|
|
"""Setup test fixtures."""
|
|
self.agent_card = create_test_agent_card()
|
|
self.mock_genai_part_converter = Mock()
|
|
self.mock_a2a_part_converter = Mock()
|
|
self.agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card=self.agent_card,
|
|
genai_part_converter=self.mock_genai_part_converter,
|
|
a2a_part_converter=self.mock_a2a_part_converter,
|
|
)
|
|
|
|
# Mock session and context
|
|
self.mock_session = Mock(spec=Session)
|
|
self.mock_session.id = "session-123"
|
|
self.mock_session.events = []
|
|
|
|
self.mock_context = Mock(spec=InvocationContext)
|
|
self.mock_context.session = self.mock_session
|
|
self.mock_context.invocation_id = "invocation-123"
|
|
self.mock_context.branch = "main"
|
|
|
|
def test_create_a2a_request_for_user_function_response_no_function_call(self):
|
|
"""Test function response request creation when no function call exists."""
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.find_matching_function_call"
|
|
) as mock_find:
|
|
mock_find.return_value = None
|
|
|
|
result = self.agent._create_a2a_request_for_user_function_response(
|
|
self.mock_context
|
|
)
|
|
|
|
assert result is None
|
|
|
|
def test_create_a2a_request_for_user_function_response_success(self):
|
|
"""Test successful function response request creation."""
|
|
# Mock function call event
|
|
mock_function_event = Mock()
|
|
mock_function_event.custom_metadata = {
|
|
A2A_METADATA_PREFIX + "task_id": "task-123"
|
|
}
|
|
mock_function_event.content = Mock()
|
|
mock_function_event.content.parts = [Mock()]
|
|
mock_function_event.get_function_calls.return_value = []
|
|
|
|
# Mock latest event with function response - set proper author
|
|
mock_latest_event = Mock()
|
|
mock_latest_event.author = "user"
|
|
self.mock_session.events = [mock_latest_event]
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.find_matching_function_call"
|
|
) as mock_find:
|
|
mock_find.return_value = mock_function_event
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.convert_event_to_a2a_message"
|
|
) as mock_convert:
|
|
# Create a proper mock A2A message
|
|
mock_a2a_message = create_autospec(A2AMessage, instance=True)
|
|
mock_a2a_message.task_id = None # Will be set by the method
|
|
mock_convert.return_value = mock_a2a_message
|
|
|
|
result = self.agent._create_a2a_request_for_user_function_response(
|
|
self.mock_context
|
|
)
|
|
|
|
assert result is not None
|
|
assert result == mock_a2a_message
|
|
assert mock_a2a_message.task_id == "task-123"
|
|
|
|
def test_construct_message_parts_from_session_success(self):
|
|
"""Test successful message parts construction from session."""
|
|
# Mock event with text content
|
|
mock_part = Mock()
|
|
mock_part.text = "Hello world"
|
|
|
|
mock_content = Mock()
|
|
mock_content.parts = [mock_part]
|
|
|
|
mock_event = Mock()
|
|
mock_event.content = mock_content
|
|
|
|
self.mock_session.events = [mock_event]
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent._present_other_agent_message"
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
mock_a2a_part = Mock()
|
|
self.mock_genai_part_converter.return_value = mock_a2a_part
|
|
|
|
parts, context_id = self.agent._construct_message_parts_from_session(
|
|
self.mock_context
|
|
)
|
|
|
|
assert len(parts) == 1
|
|
assert parts[0] == mock_a2a_part
|
|
assert context_id is None
|
|
|
|
def test_construct_message_parts_from_session_user_input_metadata(self):
|
|
"""Test that user input metadata is added for user messages."""
|
|
|
|
mock_part = Mock()
|
|
mock_content = Mock()
|
|
mock_content.parts = [mock_part]
|
|
|
|
mock_event = Mock()
|
|
mock_event.content = mock_content
|
|
mock_event.author = "user"
|
|
mock_event.custom_metadata = None
|
|
|
|
self.mock_session.events = [mock_event]
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent._present_other_agent_message"
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
# Converter returns a real A2A part; production stamps is_user_input.
|
|
a2a_part = _compat.make_text_part("hi")
|
|
self.mock_genai_part_converter.return_value = a2a_part
|
|
|
|
parts, _ = self.agent._construct_message_parts_from_session(
|
|
self.mock_context
|
|
)
|
|
|
|
assert len(parts) == 1
|
|
assert _compat.part_metadata(parts[0]).get("is_user_input") is True
|
|
|
|
def test_construct_message_parts_from_session_success_multiple_parts(self):
|
|
"""Test successful message parts construction from session."""
|
|
# Mock event with text content
|
|
mock_part = Mock()
|
|
mock_part.text = "Hello world"
|
|
|
|
mock_content = Mock()
|
|
mock_content.parts = [mock_part]
|
|
|
|
mock_event = Mock()
|
|
mock_event.content = mock_content
|
|
|
|
self.mock_session.events = [mock_event]
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent._present_other_agent_message"
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
mock_a2a_part1 = Mock()
|
|
mock_a2a_part2 = Mock()
|
|
self.mock_genai_part_converter.return_value = [
|
|
mock_a2a_part1,
|
|
mock_a2a_part2,
|
|
]
|
|
|
|
parts, context_id = self.agent._construct_message_parts_from_session(
|
|
self.mock_context
|
|
)
|
|
|
|
assert parts == [mock_a2a_part1, mock_a2a_part2]
|
|
assert context_id is None
|
|
|
|
def test_construct_message_parts_from_session_empty_events(self):
|
|
"""Test message parts construction with empty events."""
|
|
self.mock_session.events = []
|
|
|
|
parts, context_id = self.agent._construct_message_parts_from_session(
|
|
self.mock_context
|
|
)
|
|
|
|
assert parts == []
|
|
assert context_id is None
|
|
|
|
def test_construct_message_parts_from_session_stops_on_agent_reply(self):
|
|
"""Test message parts construction stops on agent reply by default."""
|
|
part1 = Mock()
|
|
part1.text = "User 1"
|
|
content1 = Mock()
|
|
content1.parts = [part1]
|
|
user1 = Mock()
|
|
user1.content = content1
|
|
user1.author = "user"
|
|
user1.custom_metadata = None
|
|
|
|
part2 = Mock()
|
|
part2.text = "Agent 1"
|
|
content2 = Mock()
|
|
content2.parts = [part2]
|
|
agent1 = Mock()
|
|
agent1.content = content2
|
|
agent1.author = self.agent.name
|
|
agent1.custom_metadata = {
|
|
A2A_METADATA_PREFIX + "response": True,
|
|
}
|
|
|
|
agent2 = Mock()
|
|
agent2.content = None
|
|
agent2.author = self.agent.name
|
|
# Just actions, no content. Not marked as a response.
|
|
agent2.actions = Mock()
|
|
agent2.custom_metadata = None
|
|
|
|
part3 = Mock()
|
|
part3.text = "User 2"
|
|
content3 = Mock()
|
|
content3.parts = [part3]
|
|
user2 = Mock()
|
|
user2.content = content3
|
|
user2.author = "user"
|
|
user2.custom_metadata = None
|
|
|
|
self.mock_session.events = [user1, agent1, user2, agent2]
|
|
|
|
def mock_converter(part):
|
|
return _compat.make_text_part(part.text)
|
|
|
|
self.mock_genai_part_converter.side_effect = mock_converter
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent._present_other_agent_message"
|
|
) as mock_present:
|
|
mock_present.side_effect = lambda event: event
|
|
parts, context_id = self.agent._construct_message_parts_from_session(
|
|
self.mock_context
|
|
)
|
|
assert len(parts) == 1
|
|
assert _compat.part_text(parts[0]) == "User 2"
|
|
assert context_id is None
|
|
|
|
def test_construct_message_parts_from_session_stateless_full_history(self):
|
|
"""Test full history for stateless agent when enabled."""
|
|
self.agent._full_history_when_stateless = True
|
|
part1 = Mock()
|
|
part1.text = "User 1"
|
|
content1 = Mock()
|
|
content1.parts = [part1]
|
|
user1 = Mock()
|
|
user1.content = content1
|
|
user1.author = "user"
|
|
user1.custom_metadata = None
|
|
|
|
part2 = Mock()
|
|
part2.text = "Agent 1"
|
|
content2 = Mock()
|
|
content2.parts = [part2]
|
|
agent1 = Mock()
|
|
agent1.content = content2
|
|
agent1.author = self.agent.name
|
|
agent1.custom_metadata = None
|
|
|
|
part3 = Mock()
|
|
part3.text = "User 2"
|
|
content3 = Mock()
|
|
content3.parts = [part3]
|
|
user2 = Mock()
|
|
user2.content = content3
|
|
user2.author = "user"
|
|
user2.custom_metadata = None
|
|
|
|
self.mock_session.events = [user1, agent1, user2]
|
|
|
|
def mock_converter(part):
|
|
return _compat.make_text_part(part.text)
|
|
|
|
self.mock_genai_part_converter.side_effect = mock_converter
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent._present_other_agent_message"
|
|
) as mock_present:
|
|
mock_present.side_effect = lambda event: event
|
|
parts, context_id = self.agent._construct_message_parts_from_session(
|
|
self.mock_context
|
|
)
|
|
assert len(parts) == 3
|
|
assert _compat.part_text(parts[0]) == "User 1"
|
|
assert _compat.part_text(parts[1]) == "Agent 1"
|
|
assert _compat.part_text(parts[2]) == "User 2"
|
|
assert context_id is None
|
|
|
|
def test_construct_message_parts_from_session_stateful_partial_history(self):
|
|
"""Test partial history for stateful agent when full history is enabled."""
|
|
self.agent._full_history_when_stateless = True
|
|
part1 = Mock()
|
|
part1.text = "User 1"
|
|
content1 = Mock()
|
|
content1.parts = [part1]
|
|
user1 = Mock()
|
|
user1.content = content1
|
|
user1.author = "user"
|
|
user1.custom_metadata = None
|
|
|
|
part2 = Mock()
|
|
part2.text = "Agent 1"
|
|
content2 = Mock()
|
|
content2.parts = [part2]
|
|
agent1 = Mock()
|
|
agent1.content = content2
|
|
agent1.author = self.agent.name
|
|
agent1.custom_metadata = {
|
|
A2A_METADATA_PREFIX + "response": True,
|
|
A2A_METADATA_PREFIX + "context_id": "ctx-1",
|
|
}
|
|
|
|
part3 = Mock()
|
|
part3.text = "User 2"
|
|
content3 = Mock()
|
|
content3.parts = [part3]
|
|
user2 = Mock()
|
|
user2.content = content3
|
|
user2.author = "user"
|
|
user2.custom_metadata = None
|
|
|
|
self.mock_session.events = [user1, agent1, user2]
|
|
|
|
def mock_converter(part):
|
|
return _compat.make_text_part(part.text)
|
|
|
|
self.mock_genai_part_converter.side_effect = mock_converter
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent._present_other_agent_message"
|
|
) as mock_present:
|
|
mock_present.side_effect = lambda event: event
|
|
parts, context_id = self.agent._construct_message_parts_from_session(
|
|
self.mock_context
|
|
)
|
|
assert len(parts) == 1
|
|
assert _compat.part_text(parts[0]) == "User 2"
|
|
assert context_id == "ctx-1"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_success_with_message(self):
|
|
"""Test successful A2A response handling with message."""
|
|
mock_a2a_message = Mock(spec=A2AMessage)
|
|
mock_a2a_message.context_id = "context-123"
|
|
|
|
# Create a proper Event mock that can handle custom_metadata
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
)
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
mock_a2a_message, self.mock_context
|
|
)
|
|
|
|
assert result == mock_event
|
|
mock_convert.assert_called_once_with(
|
|
mock_a2a_message,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.mock_a2a_part_converter,
|
|
)
|
|
# Check that metadata was added
|
|
assert result.custom_metadata is not None
|
|
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_message_converter_returns_none(self):
|
|
"""Test _handle_a2a_response returns None when message converter returns None."""
|
|
mock_a2a_message = Mock(spec=A2AMessage)
|
|
mock_a2a_message.context_id = "context-123"
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
|
|
) as mock_convert:
|
|
mock_convert.return_value = None
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
mock_a2a_message, self.mock_context
|
|
)
|
|
|
|
assert result is None
|
|
mock_convert.assert_called_once_with(
|
|
mock_a2a_message,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.mock_a2a_part_converter,
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_with_task_completed_and_no_update(self):
|
|
"""Test successful A2A response handling with non-streaming task and no update."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
mock_a2a_task.context_id = "context-123"
|
|
mock_a2a_task.status = Mock(spec=A2ATaskStatus)
|
|
mock_a2a_task.status.state = _compat.TS_COMPLETED
|
|
|
|
# Create a proper Event mock that can handle custom_metadata
|
|
mock_a2a_part = genai_types.Part.from_text(
|
|
text="test"
|
|
) # real genai part for Content
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
content=genai_types.Content(role="model", parts=[mock_a2a_part]),
|
|
)
|
|
|
|
with patch.object(
|
|
remote_a2a_agent,
|
|
"convert_a2a_task_to_event",
|
|
autospec=True,
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
(mock_a2a_task, None), self.mock_context
|
|
)
|
|
|
|
assert result == mock_event
|
|
mock_convert.assert_called_once_with(
|
|
mock_a2a_task,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.mock_a2a_part_converter,
|
|
)
|
|
# Check the parts are not updated as Thought
|
|
assert result.content.parts[0].thought is None
|
|
# Check that metadata was added
|
|
assert result.custom_metadata is not None
|
|
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
|
|
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
|
|
|
def test_construct_message_parts_from_session_preserves_order(self):
|
|
"""Test that message parts are in correct order with multi-part messages.
|
|
|
|
This test verifies the fix for the bug where _present_other_agent_message
|
|
creates multi-part messages with "For context:" prefix, and ensures the
|
|
parts are in the correct chronological order (not reversed).
|
|
"""
|
|
# Create mock events with multiple parts
|
|
# Event 1: User message
|
|
user_part = Mock()
|
|
user_part.text = "User question"
|
|
user_content = Mock()
|
|
user_content.parts = [user_part]
|
|
user_event = Mock()
|
|
user_event.content = user_content
|
|
user_event.author = "user"
|
|
|
|
# Event 2: Other agent message (will be transformed by
|
|
# _present_other_agent_message)
|
|
other_agent_part1 = Mock()
|
|
other_agent_part1.text = "For context:"
|
|
other_agent_part2 = Mock()
|
|
other_agent_part2.text = "[other_agent] said: Response text"
|
|
other_agent_content = Mock()
|
|
other_agent_content.parts = [other_agent_part1, other_agent_part2]
|
|
other_agent_event = Mock()
|
|
other_agent_event.content = other_agent_content
|
|
other_agent_event.author = "other_agent"
|
|
|
|
self.mock_session.events = [user_event, other_agent_event]
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent._present_other_agent_message"
|
|
) as mock_present:
|
|
# Mock _present_other_agent_message to return the transformed event
|
|
mock_present.return_value = other_agent_event
|
|
|
|
# Converter returns real A2A parts (production reads their metadata);
|
|
# track the conversion order for the ordering assertions below.
|
|
converted_order = []
|
|
|
|
def mock_converter(part):
|
|
converted_order.append(part.text)
|
|
return _compat.make_text_part(part.text)
|
|
|
|
self.mock_genai_part_converter.side_effect = mock_converter
|
|
|
|
parts, context_id = self.agent._construct_message_parts_from_session(
|
|
self.mock_context
|
|
)
|
|
|
|
# Verify the parts are in correct order
|
|
assert len(parts) == 3 # 1 user part + 2 other agent parts
|
|
assert context_id is None
|
|
|
|
# Verify order: user part, then "For context:", then agent message
|
|
assert converted_order[0] == "User question"
|
|
assert converted_order[1] == "For context:"
|
|
assert converted_order[2] == "[other_agent] said: Response text"
|
|
assert _compat.part_text(parts[0]) == "User question"
|
|
assert _compat.part_text(parts[1]) == "For context:"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_with_task_submitted_and_no_update(self):
|
|
"""Test successful A2A response handling with streaming task and no update."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
mock_a2a_task.context_id = "context-123"
|
|
mock_a2a_task.status = Mock(spec=A2ATaskStatus)
|
|
mock_a2a_task.status.state = _compat.TS_SUBMITTED
|
|
|
|
# Create a proper Event mock that can handle custom_metadata
|
|
mock_a2a_part = genai_types.Part.from_text(
|
|
text="test"
|
|
) # real genai part for Content
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
content=genai_types.Content(role="model", parts=[mock_a2a_part]),
|
|
)
|
|
|
|
with patch.object(
|
|
remote_a2a_agent,
|
|
"convert_a2a_task_to_event",
|
|
autospec=True,
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
(mock_a2a_task, None), self.mock_context
|
|
)
|
|
|
|
assert result == mock_event
|
|
mock_convert.assert_called_once_with(
|
|
mock_a2a_task,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.mock_a2a_part_converter,
|
|
)
|
|
# Check the parts are updated as Thought
|
|
assert result.content.parts[0].thought is True
|
|
assert result.content.parts[0].thought_signature is None
|
|
# Check that metadata was added
|
|
assert result.custom_metadata is not None
|
|
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
|
|
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"task_state,event_content",
|
|
[
|
|
pytest.param(
|
|
_compat.TS_SUBMITTED,
|
|
genai_types.Content(role="model", parts=[]),
|
|
id="submitted_empty_parts",
|
|
),
|
|
pytest.param(
|
|
_compat.TS_WORKING,
|
|
None,
|
|
id="working_no_content",
|
|
),
|
|
],
|
|
)
|
|
async def test_handle_a2a_response_with_task_missing_content(
|
|
self, task_state, event_content
|
|
):
|
|
"""Test streaming A2A response handling when content/parts are missing.
|
|
|
|
This verifies the fix for issue #3769 where the code could raise when it
|
|
tried to read parts[0] without checking for empty/missing content.
|
|
"""
|
|
mock_a2a_task = create_autospec(A2ATask, instance=True)
|
|
mock_a2a_task.id = "task-123"
|
|
mock_a2a_task.context_id = "context-123"
|
|
mock_a2a_task.status = create_autospec(A2ATaskStatus, instance=True)
|
|
mock_a2a_task.status.state = task_state
|
|
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
content=event_content,
|
|
)
|
|
|
|
with patch.object(
|
|
remote_a2a_agent,
|
|
"convert_a2a_task_to_event",
|
|
autospec=True,
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
(mock_a2a_task, None), self.mock_context
|
|
)
|
|
|
|
assert result == mock_event
|
|
assert result.custom_metadata is not None
|
|
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
|
|
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_with_task_working_and_no_update(self):
|
|
"""Test successful A2A response handling with streaming task and no update."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
mock_a2a_task.context_id = "context-123"
|
|
mock_a2a_task.status = Mock(spec=A2ATaskStatus)
|
|
mock_a2a_task.status.state = _compat.TS_WORKING
|
|
|
|
# Create a proper Event mock that can handle custom_metadata
|
|
mock_a2a_part = genai_types.Part.from_text(
|
|
text="test"
|
|
) # real genai part for Content
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
content=genai_types.Content(role="model", parts=[mock_a2a_part]),
|
|
)
|
|
|
|
with patch.object(
|
|
remote_a2a_agent,
|
|
"convert_a2a_task_to_event",
|
|
autospec=True,
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
(mock_a2a_task, None), self.mock_context
|
|
)
|
|
|
|
assert result == mock_event
|
|
mock_convert.assert_called_once_with(
|
|
mock_a2a_task,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.mock_a2a_part_converter,
|
|
)
|
|
# Check the parts are updated as Thought
|
|
assert result.content.parts[0].thought is True
|
|
assert result.content.parts[0].thought_signature is None
|
|
# Check that metadata was added
|
|
assert result.custom_metadata is not None
|
|
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
|
|
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_with_task_status_update_with_message(self):
|
|
"""Test handling of a task status update with a message."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
mock_a2a_task.context_id = "context-123"
|
|
|
|
mock_a2a_message = Mock(spec=A2AMessage)
|
|
mock_update = Mock(spec=TaskStatusUpdateEvent)
|
|
mock_update.status = Mock(A2ATaskStatus)
|
|
mock_update.status.state = _compat.TS_COMPLETED
|
|
mock_update.status.message = mock_a2a_message
|
|
|
|
# Create a proper Event mock that can handle custom_metadata
|
|
mock_a2a_part = genai_types.Part.from_text(
|
|
text="test"
|
|
) # real genai part for Content
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
content=genai_types.Content(role="model", parts=[mock_a2a_part]),
|
|
)
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
(mock_a2a_task, mock_update), self.mock_context
|
|
)
|
|
|
|
assert result == mock_event
|
|
mock_convert.assert_called_once_with(
|
|
mock_a2a_message,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.mock_a2a_part_converter,
|
|
)
|
|
# Check that metadata was added
|
|
assert result.custom_metadata is not None
|
|
assert result.content.parts[0].thought is None
|
|
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
|
|
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_with_task_status_working_update_with_message(
|
|
self,
|
|
):
|
|
"""Test handling of a task status update with a message."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
mock_a2a_task.context_id = "context-123"
|
|
|
|
mock_a2a_message = Mock(spec=A2AMessage)
|
|
mock_update = Mock(spec=TaskStatusUpdateEvent)
|
|
mock_update.status = Mock(A2ATaskStatus)
|
|
mock_update.status.state = _compat.TS_WORKING
|
|
mock_update.status.message = mock_a2a_message
|
|
|
|
# Create a proper Event mock that can handle custom_metadata
|
|
mock_a2a_part = genai_types.Part.from_text(
|
|
text="test"
|
|
) # real genai part for Content
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
content=genai_types.Content(role="model", parts=[mock_a2a_part]),
|
|
)
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
(mock_a2a_task, mock_update), self.mock_context
|
|
)
|
|
|
|
assert result == mock_event
|
|
mock_convert.assert_called_once_with(
|
|
mock_a2a_message,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.mock_a2a_part_converter,
|
|
)
|
|
# Check that metadata was added
|
|
assert result.custom_metadata is not None
|
|
assert result.content.parts[0].thought is True
|
|
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
|
|
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_with_task_status_update_no_message(self):
|
|
"""Test handling of a task status update with no message."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
|
|
mock_update = Mock(spec=TaskStatusUpdateEvent)
|
|
mock_update.status = Mock(A2ATaskStatus)
|
|
mock_update.status.state = _compat.TS_COMPLETED
|
|
mock_update.status.message = None
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
(mock_a2a_task, mock_update), self.mock_context
|
|
)
|
|
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_with_artifact_update(self):
|
|
"""Test successful A2A response handling with artifact update."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
mock_a2a_task.context_id = "context-123"
|
|
|
|
mock_artifact = Mock(spec=Artifact)
|
|
mock_update = Mock(spec=TaskArtifactUpdateEvent)
|
|
mock_update.artifact = mock_artifact
|
|
mock_update.append = False
|
|
mock_update.last_chunk = True
|
|
|
|
# Create a proper Event mock that can handle custom_metadata
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
)
|
|
|
|
with patch.object(
|
|
remote_a2a_agent,
|
|
"convert_a2a_task_to_event",
|
|
autospec=True,
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
(mock_a2a_task, mock_update), self.mock_context
|
|
)
|
|
|
|
assert result == mock_event
|
|
mock_convert.assert_called_once_with(
|
|
mock_a2a_task,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.agent._a2a_part_converter,
|
|
)
|
|
# Check that metadata was added
|
|
assert result.custom_metadata is not None
|
|
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
|
|
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_with_partial_artifact_update(self):
|
|
"""Test that partial artifact updates are ignored."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
|
|
mock_update = Mock(spec=TaskArtifactUpdateEvent)
|
|
mock_update.artifact = Mock(spec=Artifact)
|
|
mock_update.append = True
|
|
mock_update.last_chunk = False
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
(mock_a2a_task, mock_update), self.mock_context
|
|
)
|
|
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_with_real_empty_status_message(self):
|
|
"""A real status update without a message must not yield a spurious event."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
mock_a2a_task.context_id = "context-123"
|
|
|
|
# Real status update with NO message attached (empty proto Message on 1.x).
|
|
update = _compat.make_task_status_update_event(
|
|
task_id="task-123",
|
|
context_id="context-123",
|
|
status=_compat.make_task_status(_compat.TS_WORKING),
|
|
final=False,
|
|
)
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
|
|
) as mock_convert:
|
|
result = await self.agent._handle_a2a_response(
|
|
(mock_a2a_task, update), self.mock_context
|
|
)
|
|
|
|
# The empty status message must be treated as absent: the converter is not
|
|
# called and the handler produces no spurious event.
|
|
mock_convert.assert_not_called()
|
|
assert result is None
|
|
|
|
|
|
class TestRemoteA2aAgentMessageHandlingFromFactory:
|
|
"""Test message handling functionality."""
|
|
|
|
def setup_method(self):
|
|
"""Setup test fixtures."""
|
|
self.mock_a2a_part_converter = Mock()
|
|
|
|
self.agent_card = create_test_agent_card()
|
|
self.agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card=self.agent_card,
|
|
a2a_client_factory=ClientFactory(
|
|
config=ClientConfig(httpx_client=httpx.AsyncClient()),
|
|
),
|
|
a2a_part_converter=self.mock_a2a_part_converter,
|
|
)
|
|
|
|
# Mock session and context
|
|
self.mock_session = Mock(spec=Session)
|
|
self.mock_session.id = "session-123"
|
|
self.mock_session.events = []
|
|
|
|
self.mock_context = Mock(spec=InvocationContext)
|
|
self.mock_context.session = self.mock_session
|
|
self.mock_context.invocation_id = "invocation-123"
|
|
self.mock_context.branch = "main"
|
|
|
|
def test_create_a2a_request_for_user_function_response_no_function_call(self):
|
|
"""Test function response request creation when no function call exists."""
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.find_matching_function_call"
|
|
) as mock_find:
|
|
mock_find.return_value = None
|
|
|
|
result = self.agent._create_a2a_request_for_user_function_response(
|
|
self.mock_context
|
|
)
|
|
|
|
assert result is None
|
|
|
|
def test_create_a2a_request_for_user_function_response_success(self):
|
|
"""Test successful function response request creation."""
|
|
# Mock function call event
|
|
mock_function_event = Mock()
|
|
mock_function_event.custom_metadata = {
|
|
A2A_METADATA_PREFIX + "task_id": "task-123"
|
|
}
|
|
mock_function_event.content = Mock()
|
|
mock_function_event.content.parts = [Mock()]
|
|
mock_function_event.get_function_calls.return_value = []
|
|
|
|
# Mock latest event with function response - set proper author
|
|
mock_latest_event = Mock()
|
|
mock_latest_event.author = "user"
|
|
self.mock_session.events = [mock_latest_event]
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.find_matching_function_call"
|
|
) as mock_find:
|
|
mock_find.return_value = mock_function_event
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.convert_event_to_a2a_message"
|
|
) as mock_convert:
|
|
# Create a proper mock A2A message
|
|
mock_a2a_message = Mock(spec=A2AMessage)
|
|
mock_a2a_message.task_id = None # Will be set by the method
|
|
mock_convert.return_value = mock_a2a_message
|
|
|
|
result = self.agent._create_a2a_request_for_user_function_response(
|
|
self.mock_context
|
|
)
|
|
|
|
assert result is not None
|
|
assert result == mock_a2a_message
|
|
assert mock_a2a_message.task_id == "task-123"
|
|
|
|
def test_construct_message_parts_from_session_success(self):
|
|
"""Test successful message parts construction from session."""
|
|
# Mock event with text content
|
|
mock_part = Mock()
|
|
mock_part.text = "Hello world"
|
|
|
|
mock_content = Mock()
|
|
mock_content.parts = [mock_part]
|
|
|
|
mock_event = Mock()
|
|
mock_event.content = mock_content
|
|
|
|
self.mock_session.events = [mock_event]
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent._present_other_agent_message"
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
with patch.object(
|
|
self.agent, "_genai_part_converter"
|
|
) as mock_convert_part:
|
|
mock_a2a_part = Mock()
|
|
mock_convert_part.return_value = mock_a2a_part
|
|
|
|
parts, context_id = self.agent._construct_message_parts_from_session(
|
|
self.mock_context
|
|
)
|
|
|
|
assert len(parts) == 1
|
|
assert parts[0] == mock_a2a_part
|
|
assert context_id is None
|
|
|
|
def test_construct_message_parts_from_session_empty_events(self):
|
|
"""Test message parts construction with empty events."""
|
|
self.mock_session.events = []
|
|
|
|
parts, context_id = self.agent._construct_message_parts_from_session(
|
|
self.mock_context
|
|
)
|
|
|
|
assert parts == []
|
|
assert context_id is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_success_with_message(self):
|
|
"""Test successful A2A response handling with message."""
|
|
mock_a2a_message = Mock(spec=A2AMessage)
|
|
mock_a2a_message.context_id = "context-123"
|
|
|
|
# Create a proper Event mock that can handle custom_metadata
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
)
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
mock_a2a_message, self.mock_context
|
|
)
|
|
|
|
assert result == mock_event
|
|
mock_convert.assert_called_once_with(
|
|
mock_a2a_message,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.mock_a2a_part_converter,
|
|
)
|
|
# Check that metadata was added
|
|
assert result.custom_metadata is not None
|
|
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_with_task_completed_and_no_update(self):
|
|
"""Test successful A2A response handling with non-streaming task and no update."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
mock_a2a_task.context_id = "context-123"
|
|
mock_a2a_task.status = Mock(spec=A2ATaskStatus)
|
|
mock_a2a_task.status.state = _compat.TS_COMPLETED
|
|
|
|
# Create a proper Event mock that can handle custom_metadata
|
|
mock_a2a_part = genai_types.Part.from_text(
|
|
text="test"
|
|
) # real genai part for Content
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
content=genai_types.Content(role="model", parts=[mock_a2a_part]),
|
|
)
|
|
|
|
with patch.object(
|
|
remote_a2a_agent,
|
|
"convert_a2a_task_to_event",
|
|
autospec=True,
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
(mock_a2a_task, None), self.mock_context
|
|
)
|
|
|
|
assert result == mock_event
|
|
mock_convert.assert_called_once_with(
|
|
mock_a2a_task,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.mock_a2a_part_converter,
|
|
)
|
|
# Check the parts are not updated as Thought
|
|
assert result.content.parts[0].thought is None
|
|
# Check that metadata was added
|
|
assert result.custom_metadata is not None
|
|
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
|
|
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_with_task_submitted_and_no_update(self):
|
|
"""Test successful A2A response handling with streaming task and no update."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
mock_a2a_task.context_id = "context-123"
|
|
mock_a2a_task.status = Mock(spec=A2ATaskStatus)
|
|
mock_a2a_task.status.state = _compat.TS_SUBMITTED
|
|
|
|
# Create a proper Event mock that can handle custom_metadata
|
|
mock_a2a_part = genai_types.Part.from_text(
|
|
text="test"
|
|
) # real genai part for Content
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
content=genai_types.Content(role="model", parts=[mock_a2a_part]),
|
|
)
|
|
|
|
with patch.object(
|
|
remote_a2a_agent,
|
|
"convert_a2a_task_to_event",
|
|
autospec=True,
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
(mock_a2a_task, None), self.mock_context
|
|
)
|
|
|
|
assert result == mock_event
|
|
mock_convert.assert_called_once_with(
|
|
mock_a2a_task,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.agent._a2a_part_converter,
|
|
)
|
|
# Check the parts are updated as Thought
|
|
assert result.content.parts[0].thought is True
|
|
assert result.content.parts[0].thought_signature is None
|
|
# Check that metadata was added
|
|
assert result.custom_metadata is not None
|
|
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
|
|
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_with_task_status_update_with_message(self):
|
|
"""Test handling of a task status update with a message."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
mock_a2a_task.context_id = "context-123"
|
|
|
|
mock_a2a_message = Mock(spec=A2AMessage)
|
|
mock_update = Mock(spec=TaskStatusUpdateEvent)
|
|
mock_update.status = Mock(A2ATaskStatus)
|
|
mock_update.status.state = _compat.TS_COMPLETED
|
|
mock_update.status.message = mock_a2a_message
|
|
|
|
# Create a proper Event mock that can handle custom_metadata
|
|
mock_a2a_part = genai_types.Part.from_text(
|
|
text="test"
|
|
) # real genai part for Content
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
content=genai_types.Content(role="model", parts=[mock_a2a_part]),
|
|
)
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
(mock_a2a_task, mock_update), self.mock_context
|
|
)
|
|
|
|
assert result == mock_event
|
|
mock_convert.assert_called_once_with(
|
|
mock_a2a_message,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.agent._a2a_part_converter,
|
|
)
|
|
# Check that metadata was added
|
|
assert result.custom_metadata is not None
|
|
assert result.content.parts[0].thought is None
|
|
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
|
|
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_with_task_status_working_update_with_message(
|
|
self,
|
|
):
|
|
"""Test handling of a task status update with a message."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
mock_a2a_task.context_id = "context-123"
|
|
|
|
mock_a2a_message = Mock(spec=A2AMessage)
|
|
mock_update = Mock(spec=TaskStatusUpdateEvent)
|
|
mock_update.status = Mock(A2ATaskStatus)
|
|
mock_update.status.state = _compat.TS_WORKING
|
|
mock_update.status.message = mock_a2a_message
|
|
|
|
# Create a proper Event mock that can handle custom_metadata
|
|
mock_a2a_part = genai_types.Part.from_text(
|
|
text="test"
|
|
) # real genai part for Content
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
content=genai_types.Content(role="model", parts=[mock_a2a_part]),
|
|
)
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
(mock_a2a_task, mock_update), self.mock_context
|
|
)
|
|
|
|
assert result == mock_event
|
|
mock_convert.assert_called_once_with(
|
|
mock_a2a_message,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.agent._a2a_part_converter,
|
|
)
|
|
# Check that metadata was added
|
|
assert result.custom_metadata is not None
|
|
assert result.content.parts[0].thought is True
|
|
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
|
|
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_with_task_status_update_no_message(self):
|
|
"""Test handling of a task status update with no message."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
|
|
mock_update = Mock(spec=TaskStatusUpdateEvent)
|
|
mock_update.status = Mock(A2ATaskStatus)
|
|
mock_update.status.state = _compat.TS_COMPLETED
|
|
mock_update.status.message = None
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
(mock_a2a_task, mock_update), self.mock_context
|
|
)
|
|
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_with_artifact_update(self):
|
|
"""Test successful A2A response handling with artifact update."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
mock_a2a_task.context_id = "context-123"
|
|
|
|
mock_artifact = Mock(spec=Artifact)
|
|
mock_update = Mock(spec=TaskArtifactUpdateEvent)
|
|
mock_update.artifact = mock_artifact
|
|
mock_update.append = False
|
|
mock_update.last_chunk = True
|
|
|
|
# Create a proper Event mock that can handle custom_metadata
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
)
|
|
|
|
with patch.object(
|
|
remote_a2a_agent,
|
|
"convert_a2a_task_to_event",
|
|
autospec=True,
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
(mock_a2a_task, mock_update), self.mock_context
|
|
)
|
|
|
|
assert result == mock_event
|
|
mock_convert.assert_called_once_with(
|
|
mock_a2a_task,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.agent._a2a_part_converter,
|
|
)
|
|
# Check that metadata was added
|
|
assert result.custom_metadata is not None
|
|
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
|
|
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_with_partial_artifact_update(self):
|
|
"""Test that partial artifact updates are ignored."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
|
|
mock_update = Mock(spec=TaskArtifactUpdateEvent)
|
|
mock_update.artifact = Mock(spec=Artifact)
|
|
mock_update.append = True
|
|
mock_update.last_chunk = False
|
|
|
|
result = await self.agent._handle_a2a_response(
|
|
(mock_a2a_task, mock_update), self.mock_context
|
|
)
|
|
|
|
assert result is None
|
|
|
|
|
|
class TestRemoteA2aAgentMessageHandlingV2:
|
|
"""Test _handle_a2a_response_impl functionality."""
|
|
|
|
def setup_method(self):
|
|
"""Setup test fixtures."""
|
|
from google.adk.a2a.agent.config import A2aRemoteAgentConfig
|
|
|
|
self.agent_card = create_test_agent_card()
|
|
self.mock_config = Mock(spec=A2aRemoteAgentConfig)
|
|
self.mock_config.a2a_part_converter = Mock()
|
|
self.mock_config.a2a_task_converter = Mock()
|
|
self.mock_config.a2a_status_update_converter = Mock()
|
|
self.mock_config.a2a_artifact_update_converter = Mock()
|
|
self.mock_config.a2a_message_converter = Mock()
|
|
|
|
self.agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card=self.agent_card,
|
|
config=self.mock_config,
|
|
)
|
|
|
|
# Mock session and context
|
|
self.mock_session = Mock(spec=Session)
|
|
self.mock_session.id = "session-123"
|
|
self.mock_session.events = []
|
|
|
|
self.mock_context = Mock(spec=InvocationContext)
|
|
self.mock_context.session = self.mock_session
|
|
self.mock_context.invocation_id = "invocation-123"
|
|
self.mock_context.branch = "main"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_impl_with_message(self):
|
|
"""Test _handle_a2a_response_impl with A2AMessage."""
|
|
mock_a2a_message = Mock(spec=A2AMessage)
|
|
mock_a2a_message.metadata = {}
|
|
mock_a2a_message.metadata = {}
|
|
mock_a2a_message.context_id = "context-123"
|
|
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
)
|
|
self.mock_config.a2a_message_converter.return_value = mock_event
|
|
|
|
result = await self.agent._handle_a2a_response_v2(
|
|
mock_a2a_message, self.mock_context
|
|
)
|
|
|
|
assert result == mock_event
|
|
self.mock_config.a2a_message_converter.assert_called_once_with(
|
|
mock_a2a_message,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.mock_config.a2a_part_converter,
|
|
)
|
|
assert result.custom_metadata is not None
|
|
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
|
assert (
|
|
result.custom_metadata[A2A_METADATA_PREFIX + "context_id"]
|
|
== "context-123"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_impl_with_task_and_no_update(self):
|
|
"""Test _handle_a2a_response_impl with Task and no update."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
mock_a2a_task.context_id = "context-123"
|
|
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
)
|
|
self.mock_config.a2a_task_converter.return_value = mock_event
|
|
|
|
result = await self.agent._handle_a2a_response_v2(
|
|
(mock_a2a_task, None), self.mock_context
|
|
)
|
|
|
|
assert result == mock_event
|
|
self.mock_config.a2a_task_converter.assert_called_once_with(
|
|
mock_a2a_task,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.mock_config.a2a_part_converter,
|
|
)
|
|
assert result.custom_metadata is not None
|
|
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
|
|
assert result.custom_metadata[A2A_METADATA_PREFIX + "task_id"] == "task-123"
|
|
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
|
assert (
|
|
result.custom_metadata[A2A_METADATA_PREFIX + "context_id"]
|
|
== "context-123"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_impl_with_task_status_update(self):
|
|
"""Test _handle_a2a_response_impl with TaskStatusUpdateEvent."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
mock_a2a_task.context_id = None
|
|
|
|
mock_update = Mock(spec=TaskStatusUpdateEvent)
|
|
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
)
|
|
self.mock_config.a2a_status_update_converter.return_value = mock_event
|
|
|
|
result = await self.agent._handle_a2a_response_v2(
|
|
(mock_a2a_task, mock_update), self.mock_context
|
|
)
|
|
|
|
assert result == mock_event
|
|
self.mock_config.a2a_status_update_converter.assert_called_once_with(
|
|
mock_update,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.mock_config.a2a_part_converter,
|
|
)
|
|
assert result.custom_metadata is not None
|
|
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
|
|
assert result.custom_metadata[A2A_METADATA_PREFIX + "task_id"] == "task-123"
|
|
assert A2A_METADATA_PREFIX + "context_id" not in result.custom_metadata
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_impl_with_task_artifact_update(self):
|
|
"""Test _handle_a2a_response_impl with TaskArtifactUpdateEvent."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
mock_a2a_task.context_id = "context-123"
|
|
|
|
mock_update = Mock(spec=TaskArtifactUpdateEvent)
|
|
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
)
|
|
self.mock_config.a2a_artifact_update_converter.return_value = mock_event
|
|
|
|
result = await self.agent._handle_a2a_response_v2(
|
|
(mock_a2a_task, mock_update), self.mock_context
|
|
)
|
|
|
|
assert result == mock_event
|
|
self.mock_config.a2a_artifact_update_converter.assert_called_once_with(
|
|
mock_update,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.mock_config.a2a_part_converter,
|
|
)
|
|
assert result.custom_metadata is not None
|
|
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
|
|
assert result.custom_metadata[A2A_METADATA_PREFIX + "task_id"] == "task-123"
|
|
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
|
assert (
|
|
result.custom_metadata[A2A_METADATA_PREFIX + "context_id"]
|
|
== "context-123"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_impl_update_converter_returns_none(self):
|
|
"""Test _handle_a2a_response_impl when converter returns None."""
|
|
mock_a2a_task = Mock(spec=A2ATask)
|
|
mock_a2a_task.id = "task-123"
|
|
|
|
mock_update = Mock(spec=TaskArtifactUpdateEvent)
|
|
|
|
self.mock_config.a2a_artifact_update_converter.return_value = None
|
|
|
|
result = await self.agent._handle_a2a_response_v2(
|
|
(mock_a2a_task, mock_update), self.mock_context
|
|
)
|
|
|
|
assert result is None
|
|
self.mock_config.a2a_artifact_update_converter.assert_called_once_with(
|
|
mock_update,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.mock_config.a2a_part_converter,
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_impl_message_converter_returns_none(self):
|
|
"""Test _handle_a2a_response_v2 returns None when message converter returns None."""
|
|
mock_a2a_message = Mock(spec=A2AMessage)
|
|
mock_a2a_message.metadata = {}
|
|
mock_a2a_message.context_id = "context-123"
|
|
|
|
self.mock_config.a2a_message_converter.return_value = None
|
|
|
|
result = await self.agent._handle_a2a_response_v2(
|
|
mock_a2a_message, self.mock_context
|
|
)
|
|
|
|
assert result is None
|
|
self.mock_config.a2a_message_converter.assert_called_once_with(
|
|
mock_a2a_message,
|
|
self.agent.name,
|
|
self.mock_context,
|
|
self.mock_config.a2a_part_converter,
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_impl_unknown_response_type(self):
|
|
"""Test _handle_a2a_response_impl with unknown response type."""
|
|
unknown_response = object()
|
|
|
|
result = await self.agent._handle_a2a_response_v2(
|
|
unknown_response, self.mock_context
|
|
)
|
|
|
|
assert result is not None
|
|
assert result.author == self.agent.name
|
|
assert result.error_message == "Unknown A2A response type"
|
|
assert result.invocation_id == self.mock_context.invocation_id
|
|
assert result.branch == self.mock_context.branch
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_a2a_response_impl_handles_client_error(self):
|
|
"""Test _handle_a2a_response_impl catches A2AClientError."""
|
|
mock_a2a_message = Mock(spec=A2AMessage)
|
|
mock_a2a_message.metadata = {}
|
|
mock_a2a_message.metadata = {}
|
|
|
|
from google.adk.agents.remote_a2a_agent import A2AClientError
|
|
|
|
self.mock_config.a2a_message_converter.side_effect = A2AClientError(
|
|
"Test client error"
|
|
)
|
|
|
|
result = await self.agent._handle_a2a_response_v2(
|
|
mock_a2a_message, self.mock_context
|
|
)
|
|
|
|
assert result is not None
|
|
assert result.author == self.agent.name
|
|
assert (
|
|
"Failed to process A2A response: Test client error"
|
|
in result.error_message
|
|
)
|
|
assert result.invocation_id == self.mock_context.invocation_id
|
|
assert result.branch == self.mock_context.branch
|
|
|
|
|
|
class TestRemoteA2aAgentNoneConverterResults:
|
|
"""Regression tests for None converter results in both legacy and v2 handlers.
|
|
|
|
Converters can legitimately return None for messages/tasks with no convertible
|
|
parts, metadata-only events, or empty status updates. The handlers must not
|
|
crash with AttributeError when this happens.
|
|
"""
|
|
|
|
def setup_method(self):
|
|
"""Setup test fixtures."""
|
|
from google.adk.a2a.agent.config import A2aRemoteAgentConfig
|
|
|
|
self.agent_card = create_test_agent_card()
|
|
|
|
# Legacy handler agent
|
|
self.mock_a2a_part_converter = Mock()
|
|
self.legacy_agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card=self.agent_card,
|
|
a2a_part_converter=self.mock_a2a_part_converter,
|
|
)
|
|
|
|
# V2 handler agent
|
|
self.mock_config = Mock(spec=A2aRemoteAgentConfig)
|
|
self.mock_config.a2a_part_converter = Mock()
|
|
self.mock_config.a2a_task_converter = Mock()
|
|
self.mock_config.a2a_status_update_converter = Mock()
|
|
self.mock_config.a2a_artifact_update_converter = Mock()
|
|
self.mock_config.a2a_message_converter = Mock()
|
|
self.mock_config.request_interceptors = None
|
|
self.v2_agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card=self.agent_card,
|
|
config=self.mock_config,
|
|
)
|
|
|
|
# Shared mock context
|
|
self.mock_session = Mock(spec=Session)
|
|
self.mock_session.id = "session-123"
|
|
self.mock_session.events = []
|
|
|
|
self.mock_context = Mock(spec=InvocationContext)
|
|
self.mock_context.session = self.mock_session
|
|
self.mock_context.invocation_id = "invocation-123"
|
|
self.mock_context.branch = "main"
|
|
|
|
# --- V2 handler regression tests ---
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_v2_message_converter_returns_none(self):
|
|
"""V2 handler must not crash when message converter returns None."""
|
|
mock_msg = Mock(spec=A2AMessage)
|
|
mock_msg.metadata = {}
|
|
mock_msg.context_id = None
|
|
|
|
self.mock_config.a2a_message_converter.return_value = None
|
|
|
|
result = await self.v2_agent._handle_a2a_response_v2(
|
|
mock_msg, self.mock_context
|
|
)
|
|
|
|
assert result is None
|
|
self.mock_config.a2a_message_converter.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_v2_message_converter_returns_none_with_context_id(self):
|
|
"""V2 handler returns None even when message has a context_id."""
|
|
mock_msg = Mock(spec=A2AMessage)
|
|
mock_msg.metadata = {}
|
|
mock_msg.context_id = "ctx-should-not-be-accessed"
|
|
|
|
self.mock_config.a2a_message_converter.return_value = None
|
|
|
|
result = await self.v2_agent._handle_a2a_response_v2(
|
|
mock_msg, self.mock_context
|
|
)
|
|
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_v2_task_converter_returns_none(self):
|
|
"""V2 handler must not crash when task converter returns None."""
|
|
mock_task = Mock(spec=A2ATask)
|
|
mock_task.id = "task-123"
|
|
mock_task.context_id = "ctx-123"
|
|
|
|
self.mock_config.a2a_task_converter.return_value = None
|
|
|
|
result = await self.v2_agent._handle_a2a_response_v2(
|
|
(mock_task, None), self.mock_context
|
|
)
|
|
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_v2_status_update_converter_returns_none(self):
|
|
"""V2 handler must not crash when status update converter returns None."""
|
|
mock_task = Mock(spec=A2ATask)
|
|
mock_task.id = "task-123"
|
|
mock_task.context_id = None
|
|
|
|
mock_update = Mock(spec=TaskStatusUpdateEvent)
|
|
|
|
self.mock_config.a2a_status_update_converter.return_value = None
|
|
|
|
result = await self.v2_agent._handle_a2a_response_v2(
|
|
(mock_task, mock_update), self.mock_context
|
|
)
|
|
|
|
assert result is None
|
|
|
|
# --- Legacy handler regression tests ---
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_legacy_message_converter_returns_none(self):
|
|
"""Legacy handler must not crash when message converter returns None."""
|
|
mock_msg = Mock(spec=A2AMessage)
|
|
mock_msg.context_id = "context-123"
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
|
|
) as mock_convert:
|
|
mock_convert.return_value = None
|
|
|
|
result = await self.legacy_agent._handle_a2a_response(
|
|
mock_msg, self.mock_context
|
|
)
|
|
|
|
assert result is None
|
|
mock_convert.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_legacy_task_converter_returns_none_no_update(self):
|
|
"""Legacy handler must not crash when task converter returns None (no update)."""
|
|
mock_task = Mock(spec=A2ATask)
|
|
mock_task.id = "task-123"
|
|
mock_task.context_id = None
|
|
mock_task.status = Mock()
|
|
mock_task.status.state = _compat.TS_COMPLETED
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.convert_a2a_task_to_event"
|
|
) as mock_convert:
|
|
mock_convert.return_value = None
|
|
|
|
result = await self.legacy_agent._handle_a2a_response(
|
|
(mock_task, None), self.mock_context
|
|
)
|
|
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_legacy_message_converter_returns_none_status_update(self):
|
|
"""Legacy handler must not crash when message converter returns None for status update."""
|
|
mock_task = Mock(spec=A2ATask)
|
|
mock_task.id = "task-123"
|
|
mock_task.context_id = "ctx-123"
|
|
|
|
mock_update = Mock(spec=TaskStatusUpdateEvent)
|
|
mock_update.status = Mock()
|
|
mock_update.status.message = Mock()
|
|
mock_update.status.state = _compat.TS_WORKING
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
|
|
) as mock_convert:
|
|
mock_convert.return_value = None
|
|
|
|
result = await self.legacy_agent._handle_a2a_response(
|
|
(mock_task, mock_update), self.mock_context
|
|
)
|
|
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_legacy_task_converter_returns_none_artifact_update(self):
|
|
"""Legacy handler must not crash when task converter returns None for artifact update."""
|
|
mock_task = Mock(spec=A2ATask)
|
|
mock_task.id = "task-123"
|
|
mock_task.context_id = None
|
|
|
|
mock_update = Mock(spec=TaskArtifactUpdateEvent)
|
|
mock_update.append = False
|
|
mock_update.last_chunk = True
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.convert_a2a_task_to_event"
|
|
) as mock_convert:
|
|
mock_convert.return_value = None
|
|
|
|
result = await self.legacy_agent._handle_a2a_response(
|
|
(mock_task, mock_update), self.mock_context
|
|
)
|
|
|
|
assert result is None
|
|
|
|
|
|
class TestRemoteA2aAgentExecution:
|
|
"""Test agent execution functionality."""
|
|
|
|
def setup_method(self):
|
|
"""Setup test fixtures."""
|
|
self.agent_card = create_test_agent_card()
|
|
self.mock_genai_part_converter = Mock()
|
|
self.mock_a2a_part_converter = Mock()
|
|
self.agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card=self.agent_card,
|
|
genai_part_converter=self.mock_genai_part_converter,
|
|
a2a_part_converter=self.mock_a2a_part_converter,
|
|
)
|
|
|
|
# Mock session and context
|
|
self.mock_session = Mock(spec=Session)
|
|
self.mock_session.id = "session-123"
|
|
self.mock_session.events = []
|
|
self.mock_session.state = {}
|
|
|
|
self.mock_context = Mock(spec=InvocationContext)
|
|
self.mock_context.session = self.mock_session
|
|
self.mock_context.invocation_id = "invocation-123"
|
|
self.mock_context.branch = "main"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_async_impl_initialization_failure(self):
|
|
"""Test _run_async_impl when initialization fails."""
|
|
with patch.object(self.agent, "_ensure_resolved") as mock_ensure:
|
|
mock_ensure.side_effect = Exception("Initialization failed")
|
|
|
|
events = []
|
|
async for event in self.agent._run_async_impl(self.mock_context):
|
|
events.append(event)
|
|
|
|
assert len(events) == 1
|
|
assert "Failed to initialize remote A2A agent" in events[0].error_message
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_async_impl_no_message_parts(self):
|
|
"""Test _run_async_impl when no message parts are found."""
|
|
with patch.object(self.agent, "_ensure_resolved"):
|
|
with patch.object(
|
|
self.agent, "_create_a2a_request_for_user_function_response"
|
|
) as mock_create_func:
|
|
mock_create_func.return_value = None
|
|
|
|
with patch.object(
|
|
self.agent, "_construct_message_parts_from_session"
|
|
) as mock_construct:
|
|
mock_construct.return_value = (
|
|
[],
|
|
None,
|
|
) # Tuple with empty parts and no context_id
|
|
|
|
events = []
|
|
async for event in self.agent._run_async_impl(self.mock_context):
|
|
events.append(event)
|
|
|
|
assert len(events) == 1
|
|
assert events[0].content is not None
|
|
assert events[0].author == self.agent.name
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_async_impl_successful_request(self):
|
|
"""Test successful _run_async_impl execution."""
|
|
with patch.object(self.agent, "_ensure_resolved"):
|
|
with patch.object(
|
|
self.agent, "_create_a2a_request_for_user_function_response"
|
|
) as mock_create_func:
|
|
mock_create_func.return_value = None
|
|
|
|
with patch.object(
|
|
self.agent, "_construct_message_parts_from_session"
|
|
) as mock_construct:
|
|
# Use a real A2A text part so production builds a real
|
|
# A2A message that the 1.x send_message adapter can
|
|
# serialize.
|
|
mock_a2a_part = _compat.make_text_part("test")
|
|
mock_construct.return_value = (
|
|
[mock_a2a_part],
|
|
"context-123",
|
|
) # Tuple with parts and context_id
|
|
|
|
# Mock A2A client. Build the raw stream item version-
|
|
# correctly (real ``StreamResponse`` on 1.x, bare
|
|
# ``Message`` on 0.3.x) so the stream normalizer handles
|
|
# it; the dispatch itself is mocked below.
|
|
mock_a2a_client = create_autospec(spec=A2AClient, instance=True)
|
|
mock_response = _make_stream_message(
|
|
A2AMessage(
|
|
message_id="m1",
|
|
role=_compat.ROLE_USER,
|
|
parts=[mock_a2a_part],
|
|
)
|
|
)
|
|
mock_send_message = AsyncMock()
|
|
mock_send_message.__aiter__.return_value = [mock_response]
|
|
mock_a2a_client.send_message.return_value = mock_send_message
|
|
self.agent._a2a_client = mock_a2a_client
|
|
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
)
|
|
|
|
with patch.object(self.agent, "_handle_a2a_response") as mock_handle:
|
|
mock_handle.return_value = mock_event
|
|
|
|
# Mock the logging functions to avoid iteration issues
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.build_a2a_request_log"
|
|
) as mock_req_log:
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.build_a2a_response_log"
|
|
) as mock_resp_log:
|
|
mock_req_log.return_value = "Mock request log"
|
|
mock_resp_log.return_value = "Mock response log"
|
|
|
|
# Patch the production serializer so metadata
|
|
# stamping does not run MessageToDict on the
|
|
# mock response (which crashes on 1.x).
|
|
with patch(
|
|
"google.adk.a2a._compat.a2a_to_dict",
|
|
return_value={"k": "v"},
|
|
):
|
|
# Execute
|
|
events = []
|
|
async for event in self.agent._run_async_impl(
|
|
self.mock_context
|
|
):
|
|
events.append(event)
|
|
|
|
assert len(events) == 1
|
|
assert events[0] == mock_event
|
|
assert (
|
|
A2A_METADATA_PREFIX + "request"
|
|
in mock_event.custom_metadata
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_async_impl_a2a_client_error(self):
|
|
"""Test _run_async_impl when A2A send_message fails."""
|
|
with patch.object(self.agent, "_ensure_resolved"):
|
|
with patch.object(
|
|
self.agent, "_create_a2a_request_for_user_function_response"
|
|
) as mock_create_func:
|
|
mock_create_func.return_value = None
|
|
|
|
with patch.object(
|
|
self.agent, "_construct_message_parts_from_session"
|
|
) as mock_construct:
|
|
# Return a real A2A part so the request Message builds and can
|
|
# be serialized in the error path on both SDK versions.
|
|
mock_a2a_part = _compat.make_text_part("test")
|
|
mock_construct.return_value = (
|
|
[mock_a2a_part],
|
|
"context-123",
|
|
) # Tuple with parts and context_id
|
|
|
|
# Mock A2A client that throws an exception
|
|
mock_a2a_client = AsyncMock()
|
|
mock_a2a_client.send_message.side_effect = Exception("Send failed")
|
|
self.agent._a2a_client = mock_a2a_client
|
|
|
|
# Mock the logging functions to avoid iteration issues
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.build_a2a_request_log"
|
|
) as mock_req_log:
|
|
mock_req_log.return_value = "Mock request log"
|
|
|
|
events = []
|
|
async for event in self.agent._run_async_impl(self.mock_context):
|
|
events.append(event)
|
|
|
|
assert len(events) == 1
|
|
assert "A2A request failed" in events[0].error_message
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_live_impl_not_implemented(self):
|
|
"""Test that _run_live_impl raises NotImplementedError."""
|
|
with pytest.raises(
|
|
NotImplementedError, match="_run_live_impl.*not implemented"
|
|
):
|
|
async for _ in self.agent._run_live_impl(self.mock_context):
|
|
pass
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_async_impl_with_meta_provider(self):
|
|
"""Test _run_async_impl with a2a_request_meta_provider."""
|
|
mock_meta_provider = Mock()
|
|
request_metadata = {"custom_meta": "value"}
|
|
mock_meta_provider.return_value = request_metadata
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card=self.agent_card,
|
|
genai_part_converter=self.mock_genai_part_converter,
|
|
a2a_part_converter=self.mock_a2a_part_converter,
|
|
a2a_request_meta_provider=mock_meta_provider,
|
|
)
|
|
|
|
with patch.object(agent, "_ensure_resolved"):
|
|
with patch.object(
|
|
agent, "_create_a2a_request_for_user_function_response"
|
|
) as mock_create_func:
|
|
mock_create_func.return_value = None
|
|
|
|
with patch.object(
|
|
agent, "_construct_message_parts_from_session"
|
|
) as mock_construct:
|
|
# Use a real A2A text part so production builds a real
|
|
# A2A message that the 1.x send_message adapter can
|
|
# serialize (it CopyFrom()s the message into a proto).
|
|
mock_a2a_part = _compat.make_text_part("test")
|
|
mock_construct.return_value = (
|
|
[mock_a2a_part],
|
|
"context-123",
|
|
) # Tuple with parts and context_id
|
|
|
|
# Mock A2A client. The raw stream item is built
|
|
# version-correctly (a real ``StreamResponse`` on 1.x, a
|
|
# bare ``Message`` on 0.3.x) so production's
|
|
# the stream normalizer handles it; the dispatch itself
|
|
# is mocked via ``_handle_a2a_response`` below.
|
|
mock_a2a_client = create_autospec(spec=A2AClient, instance=True)
|
|
mock_response = _make_stream_message(
|
|
A2AMessage(
|
|
message_id="m1",
|
|
role=_compat.ROLE_USER,
|
|
parts=[mock_a2a_part],
|
|
)
|
|
)
|
|
mock_send_message = AsyncMock()
|
|
mock_send_message.__aiter__.return_value = [mock_response]
|
|
mock_a2a_client.send_message.return_value = mock_send_message
|
|
# Use the locally-created ``agent`` (the one with the
|
|
# meta_provider and the patched _create/_construct), not
|
|
# ``self.agent``.
|
|
agent._a2a_client = mock_a2a_client
|
|
|
|
mock_event = Event(
|
|
author=agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
)
|
|
|
|
with patch.object(agent, "_handle_a2a_response") as mock_handle:
|
|
mock_handle.return_value = mock_event
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.build_a2a_request_log"
|
|
) as mock_req_log:
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.build_a2a_response_log"
|
|
) as mock_resp_log:
|
|
mock_req_log.return_value = "Mock request log"
|
|
mock_resp_log.return_value = "Mock response log"
|
|
|
|
with patch(
|
|
"google.adk.a2a._compat.a2a_to_dict",
|
|
return_value={"k": "v"},
|
|
):
|
|
events = []
|
|
async for event in agent._run_async_impl(self.mock_context):
|
|
events.append(event)
|
|
|
|
assert len(events) == 1
|
|
assert events[0] == mock_event
|
|
assert (
|
|
A2A_METADATA_PREFIX + "request"
|
|
in mock_event.custom_metadata
|
|
)
|
|
|
|
|
|
class TestRemoteA2aAgentExecutionFromFactory:
|
|
"""Test agent execution functionality (factory-constructed client)."""
|
|
|
|
def setup_method(self):
|
|
"""Setup test fixtures."""
|
|
self.agent_card = create_test_agent_card()
|
|
self.agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card=self.agent_card,
|
|
a2a_client_factory=ClientFactory(
|
|
config=ClientConfig(httpx_client=httpx.AsyncClient()),
|
|
),
|
|
)
|
|
|
|
# Mock session and context
|
|
self.mock_session = Mock(spec=Session)
|
|
self.mock_session.id = "session-123"
|
|
self.mock_session.events = []
|
|
self.mock_session.state = {}
|
|
|
|
self.mock_context = Mock(spec=InvocationContext)
|
|
self.mock_context.session = self.mock_session
|
|
self.mock_context.invocation_id = "invocation-123"
|
|
self.mock_context.branch = "main"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_async_impl_initialization_failure(self):
|
|
"""Test _run_async_impl when initialization fails."""
|
|
with patch.object(self.agent, "_ensure_resolved") as mock_ensure:
|
|
mock_ensure.side_effect = Exception("Initialization failed")
|
|
|
|
events = []
|
|
async for event in self.agent._run_async_impl(self.mock_context):
|
|
events.append(event)
|
|
|
|
assert len(events) == 1
|
|
assert "Failed to initialize remote A2A agent" in events[0].error_message
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_async_impl_no_message_parts(self):
|
|
"""Test _run_async_impl when no message parts are found."""
|
|
with patch.object(self.agent, "_ensure_resolved"):
|
|
with patch.object(
|
|
self.agent, "_create_a2a_request_for_user_function_response"
|
|
) as mock_create_func:
|
|
mock_create_func.return_value = None
|
|
|
|
with patch.object(
|
|
self.agent, "_construct_message_parts_from_session"
|
|
) as mock_construct:
|
|
mock_construct.return_value = (
|
|
[],
|
|
None,
|
|
) # Tuple with empty parts and no context_id
|
|
|
|
events = []
|
|
async for event in self.agent._run_async_impl(self.mock_context):
|
|
events.append(event)
|
|
|
|
assert len(events) == 1
|
|
assert events[0].content is not None
|
|
assert events[0].author == self.agent.name
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_async_impl_successful_request(self):
|
|
"""Test successful _run_async_impl execution."""
|
|
with patch.object(self.agent, "_ensure_resolved"):
|
|
with patch.object(
|
|
self.agent, "_create_a2a_request_for_user_function_response"
|
|
) as mock_create_func:
|
|
mock_create_func.return_value = None
|
|
|
|
with patch.object(
|
|
self.agent, "_construct_message_parts_from_session"
|
|
) as mock_construct:
|
|
# Use a real A2A text part so production builds a real
|
|
# A2A message that the 1.x send_message adapter can
|
|
# serialize.
|
|
mock_a2a_part = _compat.make_text_part("test")
|
|
mock_construct.return_value = (
|
|
[mock_a2a_part],
|
|
"context-123",
|
|
) # Tuple with parts and context_id
|
|
|
|
# Build the raw stream item version-correctly (real
|
|
# StreamResponse on 1.x, bare Message on 0.3.x) so
|
|
# the stream normalizer handles it; the dispatch is mocked.
|
|
mock_a2a_client = create_autospec(spec=A2AClient, instance=True)
|
|
mock_response = _make_stream_message(
|
|
A2AMessage(
|
|
message_id="m1",
|
|
role=_compat.ROLE_USER,
|
|
parts=[mock_a2a_part],
|
|
)
|
|
)
|
|
mock_send_message = AsyncMock()
|
|
mock_send_message.__aiter__.return_value = [mock_response]
|
|
mock_a2a_client.send_message.return_value = mock_send_message
|
|
self.agent._a2a_client = mock_a2a_client
|
|
|
|
mock_event = Event(
|
|
author=self.agent.name,
|
|
invocation_id=self.mock_context.invocation_id,
|
|
branch=self.mock_context.branch,
|
|
)
|
|
|
|
with patch.object(self.agent, "_handle_a2a_response") as mock_handle:
|
|
mock_handle.return_value = mock_event
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.build_a2a_request_log"
|
|
) as mock_req_log:
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.build_a2a_response_log"
|
|
) as mock_resp_log:
|
|
mock_req_log.return_value = "Mock request log"
|
|
mock_resp_log.return_value = "Mock response log"
|
|
|
|
with patch(
|
|
"google.adk.a2a._compat.a2a_to_dict",
|
|
return_value={"k": "v"},
|
|
):
|
|
events = []
|
|
async for event in self.agent._run_async_impl(
|
|
self.mock_context
|
|
):
|
|
events.append(event)
|
|
|
|
assert len(events) == 1
|
|
assert events[0] == mock_event
|
|
assert (
|
|
A2A_METADATA_PREFIX + "request"
|
|
in mock_event.custom_metadata
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_async_impl_a2a_client_error(self):
|
|
"""Test _run_async_impl when A2A send_message fails."""
|
|
with patch.object(self.agent, "_ensure_resolved"):
|
|
with patch.object(
|
|
self.agent, "_create_a2a_request_for_user_function_response"
|
|
) as mock_create_func:
|
|
mock_create_func.return_value = None
|
|
|
|
with patch.object(
|
|
self.agent, "_construct_message_parts_from_session"
|
|
) as mock_construct:
|
|
# Return a real A2A part so the request Message builds and can
|
|
# be serialized in the error path on both SDK versions.
|
|
mock_a2a_part = _compat.make_text_part("test")
|
|
mock_construct.return_value = (
|
|
[mock_a2a_part],
|
|
"context-123",
|
|
) # Tuple with parts and context_id
|
|
|
|
# Mock A2A client that throws an exception
|
|
mock_a2a_client = AsyncMock()
|
|
mock_a2a_client.send_message.side_effect = Exception("Send failed")
|
|
self.agent._a2a_client = mock_a2a_client
|
|
|
|
# Mock the logging functions to avoid iteration issues
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.build_a2a_request_log"
|
|
) as mock_req_log:
|
|
mock_req_log.return_value = "Mock request log"
|
|
|
|
events = []
|
|
async for event in self.agent._run_async_impl(self.mock_context):
|
|
events.append(event)
|
|
|
|
assert len(events) == 1
|
|
assert "A2A request failed" in events[0].error_message
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_live_impl_not_implemented(self):
|
|
"""Test that _run_live_impl raises NotImplementedError."""
|
|
with pytest.raises(
|
|
NotImplementedError, match="_run_live_impl.*not implemented"
|
|
):
|
|
async for _ in self.agent._run_live_impl(self.mock_context):
|
|
pass
|
|
|
|
|
|
class TestRemoteA2aAgentCleanup:
|
|
"""Test cleanup functionality."""
|
|
|
|
def setup_method(self):
|
|
"""Setup test fixtures."""
|
|
self.agent_card = create_test_agent_card()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_owns_httpx_client(self):
|
|
"""Test cleanup when agent owns httpx client."""
|
|
agent = RemoteA2aAgent(name="test_agent", agent_card=self.agent_card)
|
|
|
|
# Set up owned client
|
|
mock_client = AsyncMock()
|
|
agent._httpx_client = mock_client
|
|
agent._httpx_client_needs_cleanup = True
|
|
|
|
await agent.cleanup()
|
|
|
|
mock_client.aclose.assert_called_once()
|
|
assert agent._httpx_client is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_owns_httpx_client_factory(self):
|
|
"""Test cleanup when agent owns httpx client."""
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card=self.agent_card,
|
|
a2a_client_factory=ClientFactory(config=ClientConfig()),
|
|
)
|
|
|
|
# Set up owned client
|
|
mock_client = AsyncMock()
|
|
agent._httpx_client = mock_client
|
|
agent._httpx_client_needs_cleanup = True
|
|
|
|
await agent.cleanup()
|
|
|
|
mock_client.aclose.assert_called_once()
|
|
assert agent._httpx_client is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_does_not_own_httpx_client(self):
|
|
"""Test cleanup when agent does not own httpx client."""
|
|
shared_client = AsyncMock()
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card=self.agent_card,
|
|
httpx_client=shared_client,
|
|
)
|
|
|
|
await agent.cleanup()
|
|
|
|
# Should not close shared client
|
|
shared_client.aclose.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_does_not_own_httpx_client_factory(self):
|
|
"""Test cleanup when agent does not own httpx client."""
|
|
shared_client = AsyncMock()
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card=self.agent_card,
|
|
a2a_client_factory=ClientFactory(
|
|
config=ClientConfig(httpx_client=shared_client)
|
|
),
|
|
)
|
|
|
|
await agent.cleanup()
|
|
|
|
# Should not close shared client
|
|
shared_client.aclose.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_client_close_error(self):
|
|
"""Test cleanup when client close raises error."""
|
|
agent = RemoteA2aAgent(name="test_agent", agent_card=self.agent_card)
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.aclose.side_effect = Exception("Close failed")
|
|
agent._httpx_client = mock_client
|
|
agent._httpx_client_needs_cleanup = True
|
|
|
|
# Should not raise exception
|
|
await agent.cleanup()
|
|
assert agent._httpx_client is None
|
|
|
|
|
|
class TestRemoteA2aAgentIntegration:
|
|
"""Integration tests for RemoteA2aAgent."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_full_workflow_with_direct_agent_card(self):
|
|
"""Test full workflow with direct agent card."""
|
|
agent_card = create_test_agent_card()
|
|
|
|
agent = RemoteA2aAgent(name="test_agent", agent_card=agent_card)
|
|
|
|
# Use a real genai Part for the session content. The instance's
|
|
# part converter is bound at construction to the real
|
|
# convert_genai_part_to_a2a_part (the module-level patch below does
|
|
# not rebind it), so a real Part is needed to produce a serializable
|
|
# A2A part on both SDK versions.
|
|
mock_part = genai_types.Part.from_text(text="Hello world")
|
|
|
|
mock_content = Mock()
|
|
mock_content.parts = [mock_part]
|
|
|
|
mock_event = Mock()
|
|
mock_event.content = mock_content
|
|
|
|
mock_session = Mock(spec=Session)
|
|
mock_session.id = "session-123"
|
|
mock_session.events = [mock_event]
|
|
mock_session.state = {}
|
|
|
|
mock_context = Mock(spec=InvocationContext)
|
|
mock_context.session = mock_session
|
|
mock_context.invocation_id = "invocation-123"
|
|
mock_context.branch = "main"
|
|
|
|
# Mock dependencies
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent._present_other_agent_message"
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.convert_genai_part_to_a2a_part"
|
|
) as mock_convert_part:
|
|
# Return a real A2A text part so production builds a real
|
|
# A2A message that the 1.x send_message adapter can serialize.
|
|
mock_a2a_part = _compat.make_text_part("test")
|
|
mock_convert_part.return_value = mock_a2a_part
|
|
|
|
with patch("httpx.AsyncClient") as mock_httpx_client_class:
|
|
mock_httpx_client = AsyncMock()
|
|
mock_httpx_client_class.return_value = mock_httpx_client
|
|
|
|
with patch.object(agent, "_a2a_client") as mock_a2a_client:
|
|
# Build a real message (wrapped in a StreamResponse on
|
|
# 1.x) so production's stream normalizer /
|
|
# dispatch treat it as a message; the conversion itself
|
|
# is mocked below.
|
|
mock_a2a_message = A2AMessage(
|
|
message_id="m1",
|
|
role=_compat.ROLE_USER,
|
|
parts=[mock_a2a_part],
|
|
context_id="context-123",
|
|
)
|
|
mock_response = _make_stream_message(mock_a2a_message)
|
|
|
|
mock_send_message = AsyncMock()
|
|
mock_send_message.__aiter__.return_value = [mock_response]
|
|
mock_a2a_client.send_message.return_value = mock_send_message
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
|
|
) as mock_convert_event:
|
|
mock_result_event = Event(
|
|
author=agent.name,
|
|
invocation_id=mock_context.invocation_id,
|
|
branch=mock_context.branch,
|
|
)
|
|
mock_convert_event.return_value = mock_result_event
|
|
|
|
# Mock the logging functions to avoid iteration issues
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.build_a2a_request_log"
|
|
) as mock_req_log:
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.build_a2a_response_log"
|
|
) as mock_resp_log:
|
|
mock_req_log.return_value = "Mock request log"
|
|
mock_resp_log.return_value = "Mock response log"
|
|
|
|
# Patch the production serializer so metadata
|
|
# stamping does not run MessageToDict on the
|
|
# mock response (which crashes on 1.x).
|
|
with patch(
|
|
"google.adk.a2a._compat.a2a_to_dict",
|
|
return_value={"k": "v"},
|
|
):
|
|
# Execute
|
|
events = []
|
|
async for event in agent._run_async_impl(mock_context):
|
|
events.append(event)
|
|
|
|
assert len(events) == 1
|
|
assert events[0] == mock_result_event
|
|
assert (
|
|
A2A_METADATA_PREFIX + "request"
|
|
in mock_result_event.custom_metadata
|
|
)
|
|
|
|
# Verify A2A client was called
|
|
mock_a2a_client.send_message.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_full_workflow_with_direct_agent_card_and_factory(self):
|
|
"""Test full workflow with direct agent card."""
|
|
agent_card = create_test_agent_card()
|
|
|
|
agent = RemoteA2aAgent(
|
|
name="test_agent",
|
|
agent_card=agent_card,
|
|
a2a_client_factory=ClientFactory(config=ClientConfig()),
|
|
)
|
|
|
|
# Use a real genai Part for the session content. The instance's
|
|
# part converter is bound at construction to the real
|
|
# convert_genai_part_to_a2a_part (the module-level patch below does
|
|
# not rebind it), so a real Part is needed to produce a serializable
|
|
# A2A part on both SDK versions.
|
|
mock_part = genai_types.Part.from_text(text="Hello world")
|
|
|
|
mock_content = Mock()
|
|
mock_content.parts = [mock_part]
|
|
|
|
mock_event = Mock()
|
|
mock_event.content = mock_content
|
|
|
|
mock_session = Mock(spec=Session)
|
|
mock_session.id = "session-123"
|
|
mock_session.events = [mock_event]
|
|
mock_session.state = {}
|
|
|
|
mock_context = Mock(spec=InvocationContext)
|
|
mock_context.session = mock_session
|
|
mock_context.invocation_id = "invocation-123"
|
|
mock_context.branch = "main"
|
|
|
|
# Mock dependencies
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent._present_other_agent_message"
|
|
) as mock_convert:
|
|
mock_convert.return_value = mock_event
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.convert_genai_part_to_a2a_part"
|
|
) as mock_convert_part:
|
|
# Return a real A2A text part so production builds a real
|
|
# A2A message that the 1.x send_message adapter can serialize.
|
|
mock_a2a_part = _compat.make_text_part("test")
|
|
mock_convert_part.return_value = mock_a2a_part
|
|
|
|
with patch("httpx.AsyncClient") as mock_httpx_client_class:
|
|
mock_httpx_client = AsyncMock()
|
|
mock_httpx_client_class.return_value = mock_httpx_client
|
|
|
|
with patch.object(agent, "_a2a_client") as mock_a2a_client:
|
|
# Build a real message (wrapped in a StreamResponse on
|
|
# 1.x) so production's stream normalizer /
|
|
# dispatch treat it as a message; the conversion itself
|
|
# is mocked below.
|
|
mock_a2a_message = A2AMessage(
|
|
message_id="m1",
|
|
role=_compat.ROLE_USER,
|
|
parts=[mock_a2a_part],
|
|
context_id="context-123",
|
|
)
|
|
mock_response = _make_stream_message(mock_a2a_message)
|
|
|
|
mock_send_message = AsyncMock()
|
|
mock_send_message.__aiter__.return_value = [mock_response]
|
|
mock_a2a_client.send_message.return_value = mock_send_message
|
|
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
|
|
) as mock_convert_event:
|
|
mock_result_event = Event(
|
|
author=agent.name,
|
|
invocation_id=mock_context.invocation_id,
|
|
branch=mock_context.branch,
|
|
)
|
|
mock_convert_event.return_value = mock_result_event
|
|
|
|
# Mock the logging functions to avoid iteration issues
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.build_a2a_request_log"
|
|
) as mock_req_log:
|
|
with patch(
|
|
"google.adk.agents.remote_a2a_agent.build_a2a_response_log"
|
|
) as mock_resp_log:
|
|
mock_req_log.return_value = "Mock request log"
|
|
mock_resp_log.return_value = "Mock response log"
|
|
|
|
# Patch the production serializer so metadata
|
|
# stamping does not run MessageToDict on the
|
|
# mock response (which crashes on 1.x).
|
|
with patch(
|
|
"google.adk.a2a._compat.a2a_to_dict",
|
|
return_value={"k": "v"},
|
|
):
|
|
# Execute
|
|
events = []
|
|
async for event in agent._run_async_impl(mock_context):
|
|
events.append(event)
|
|
|
|
assert len(events) == 1
|
|
assert events[0] == mock_result_event
|
|
assert (
|
|
A2A_METADATA_PREFIX + "request"
|
|
in mock_result_event.custom_metadata
|
|
)
|
|
|
|
# Verify A2A client was called
|
|
mock_a2a_client.send_message.assert_called_once()
|
|
|
|
|
|
class TestRemoteA2aAgentInterceptors:
|
|
|
|
@pytest.fixture
|
|
def mock_context(self):
|
|
ctx = Mock(spec=InvocationContext)
|
|
ctx.session = Mock()
|
|
ctx.session.state = {"key": "value"}
|
|
return ctx
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_before_request_interceptors_none(self, mock_context):
|
|
request = Mock(spec=A2AMessage)
|
|
result_req, params = await execute_before_request_interceptors(
|
|
None, mock_context, request
|
|
)
|
|
assert result_req is request
|
|
assert params.client_call_context.state == {"key": "value"}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_before_request_interceptors_empty(self, mock_context):
|
|
request = Mock(spec=A2AMessage)
|
|
result_req, params = await execute_before_request_interceptors(
|
|
[], mock_context, request
|
|
)
|
|
assert result_req is request
|
|
assert params.client_call_context.state == {"key": "value"}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_before_request_interceptors_success(
|
|
self, mock_context
|
|
):
|
|
request = Mock(spec=A2AMessage)
|
|
new_request = Mock(spec=A2AMessage)
|
|
|
|
interceptor1 = Mock(spec=RequestInterceptor)
|
|
interceptor1.before_request = AsyncMock(
|
|
return_value=(
|
|
new_request,
|
|
ParametersConfig(
|
|
client_call_context=_compat.ClientCallContext(
|
|
state={"updated": "true"}
|
|
)
|
|
),
|
|
)
|
|
)
|
|
|
|
result_req, params = await execute_before_request_interceptors(
|
|
[interceptor1], mock_context, request
|
|
)
|
|
|
|
assert result_req is new_request
|
|
assert params.client_call_context.state == {"updated": "true"}
|
|
interceptor1.before_request.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_before_request_interceptors_returns_event(
|
|
self, mock_context
|
|
):
|
|
request = Mock(spec=A2AMessage)
|
|
event = Mock(spec=Event)
|
|
|
|
interceptor1 = Mock(spec=RequestInterceptor)
|
|
interceptor1.before_request = AsyncMock(
|
|
return_value=(
|
|
event,
|
|
ParametersConfig(
|
|
client_call_context=_compat.ClientCallContext(
|
|
state={"updated": "true"}
|
|
)
|
|
),
|
|
)
|
|
)
|
|
|
|
interceptor2 = Mock(spec=RequestInterceptor)
|
|
interceptor2.before_request = AsyncMock()
|
|
|
|
result, params = await execute_before_request_interceptors(
|
|
[interceptor1, interceptor2], mock_context, request
|
|
)
|
|
|
|
assert result is event
|
|
assert params.client_call_context.state == {"updated": "true"}
|
|
interceptor1.before_request.assert_called_once()
|
|
interceptor2.before_request.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_before_request_interceptors_no_before_request(
|
|
self, mock_context
|
|
):
|
|
request = Mock(spec=A2AMessage)
|
|
|
|
interceptor1 = Mock(spec=RequestInterceptor)
|
|
interceptor1.before_request = None
|
|
|
|
result_req, params = await execute_before_request_interceptors(
|
|
[interceptor1], mock_context, request
|
|
)
|
|
|
|
assert result_req is request
|
|
assert params.client_call_context.state == {"key": "value"}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_after_request_interceptors_none(self, mock_context):
|
|
response = Mock(spec=A2AMessage)
|
|
event = Mock(spec=Event)
|
|
result = await execute_after_request_interceptors(
|
|
None, mock_context, response, event
|
|
)
|
|
assert result is event
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_after_request_interceptors_empty(self, mock_context):
|
|
response = Mock(spec=A2AMessage)
|
|
event = Mock(spec=Event)
|
|
result = await execute_after_request_interceptors(
|
|
[], mock_context, response, event
|
|
)
|
|
assert result is event
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_after_request_interceptors_success(self, mock_context):
|
|
response = Mock(spec=A2AMessage)
|
|
event = Mock(spec=Event)
|
|
new_event = Mock(spec=Event)
|
|
|
|
interceptor1 = Mock(spec=RequestInterceptor)
|
|
interceptor1.after_request = AsyncMock(return_value=new_event)
|
|
|
|
result = await execute_after_request_interceptors(
|
|
[interceptor1], mock_context, response, event
|
|
)
|
|
|
|
assert result is new_event
|
|
interceptor1.after_request.assert_called_once_with(
|
|
mock_context, response, event
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_after_request_interceptors_reverse_order(
|
|
self, mock_context
|
|
):
|
|
response = Mock(spec=A2AMessage)
|
|
event = Mock(spec=Event)
|
|
event1 = Mock(spec=Event)
|
|
event2 = Mock(spec=Event)
|
|
|
|
interceptor1 = Mock(spec=RequestInterceptor)
|
|
interceptor1.after_request = AsyncMock(return_value=event1)
|
|
|
|
interceptor2 = Mock(spec=RequestInterceptor)
|
|
interceptor2.after_request = AsyncMock(return_value=event2)
|
|
|
|
result = await execute_after_request_interceptors(
|
|
[interceptor1, interceptor2], mock_context, response, event
|
|
)
|
|
|
|
assert result is event1
|
|
interceptor2.after_request.assert_called_once_with(
|
|
mock_context, response, event
|
|
)
|
|
interceptor1.after_request.assert_called_once_with(
|
|
mock_context, response, event2
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_after_request_interceptors_returns_none(
|
|
self, mock_context
|
|
):
|
|
response = Mock(spec=A2AMessage)
|
|
event = Mock(spec=Event)
|
|
|
|
interceptor1 = Mock(spec=RequestInterceptor)
|
|
interceptor1.after_request = AsyncMock()
|
|
|
|
interceptor2 = Mock(spec=RequestInterceptor)
|
|
interceptor2.after_request = AsyncMock(return_value=None)
|
|
|
|
result = await execute_after_request_interceptors(
|
|
[interceptor1, interceptor2], mock_context, response, event
|
|
)
|
|
|
|
assert result is None
|
|
interceptor2.after_request.assert_called_once_with(
|
|
mock_context, response, event
|
|
)
|
|
interceptor1.after_request.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_after_request_interceptors_no_after_request(
|
|
self, mock_context
|
|
):
|
|
response = Mock(spec=A2AMessage)
|
|
event = Mock(spec=Event)
|
|
|
|
interceptor1 = Mock(spec=RequestInterceptor)
|
|
interceptor1.after_request = None
|
|
|
|
result = await execute_after_request_interceptors(
|
|
[interceptor1], mock_context, response, event
|
|
)
|
|
|
|
assert result is event
|
|
|
|
|
|
class TestRemoteA2aAgentDeepcopy:
|
|
"""Test deepcopy functionality for RemoteA2aAgent and its config."""
|
|
|
|
def test_deepcopy_config(self):
|
|
"""Test that A2aRemoteAgentConfig can be deepcopied with interceptors."""
|
|
config = A2aRemoteAgentConfig()
|
|
mock_interceptor = Mock()
|
|
config.request_interceptors = [mock_interceptor]
|
|
|
|
copied_config = copy.deepcopy(config)
|
|
assert copied_config is not None
|
|
|
|
# Verify that functions are shared (by reference)
|
|
assert copied_config.a2a_message_converter is config.a2a_message_converter
|
|
|
|
# Verify that request_interceptors list was copied
|
|
assert copied_config.request_interceptors is not None
|
|
assert len(copied_config.request_interceptors) == 1
|
|
# Standard objects inside lists should be deepcopied (new instances)
|
|
assert (
|
|
copied_config.request_interceptors[0]
|
|
is not config.request_interceptors[0]
|
|
)
|