Skip to content

Commit ce49323

Browse files
caohy1988claude
andcommitted
test(plugins): add e2e callback tests for TRANSFER_A2A in emitted rows
Exercises before_tool_callback and after_tool_callback with a real TransferToAgentTool and asserts that tool_origin="TRANSFER_A2A" appears in the serialized BigQuery row content, closing the testing gap between the unit-level _resolve_transfer_origin tests and the actual callback wiring. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 73aba76 commit ce49323

File tree

1 file changed

+67
-0
lines changed

1 file changed

+67
-0
lines changed

tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4437,6 +4437,73 @@ def test_transfer_tool_unknown_name_returns_transfer_agent(self):
44374437
)
44384438
assert result == "TRANSFER_AGENT"
44394439

4440+
@pytest.mark.asyncio
4441+
async def test_before_tool_callback_emits_transfer_a2a_in_row(
4442+
self,
4443+
bq_plugin_inst,
4444+
mock_write_client,
4445+
tool_context,
4446+
dummy_arrow_schema,
4447+
):
4448+
"""End-to-end: before_tool_callback emits tool_origin=TRANSFER_A2A."""
4449+
from google.adk.tools.transfer_to_agent_tool import TransferToAgentTool
4450+
4451+
tool = TransferToAgentTool(
4452+
agent_names=["local", "remote"],
4453+
target_origin_by_name={
4454+
"local": "TRANSFER_AGENT",
4455+
"remote": "TRANSFER_A2A",
4456+
},
4457+
)
4458+
bigquery_agent_analytics_plugin.TraceManager.push_span(tool_context)
4459+
await bq_plugin_inst.before_tool_callback(
4460+
tool=tool,
4461+
tool_args={"agent_name": "remote"},
4462+
tool_context=tool_context,
4463+
)
4464+
await asyncio.sleep(0.01)
4465+
log_entry = await _get_captured_event_dict_async(
4466+
mock_write_client, dummy_arrow_schema
4467+
)
4468+
_assert_common_fields(log_entry, "TOOL_STARTING")
4469+
content_dict = json.loads(log_entry["content"])
4470+
assert content_dict["tool"] == "transfer_to_agent"
4471+
assert content_dict["tool_origin"] == "TRANSFER_A2A"
4472+
4473+
@pytest.mark.asyncio
4474+
async def test_after_tool_callback_emits_transfer_a2a_in_row(
4475+
self,
4476+
bq_plugin_inst,
4477+
mock_write_client,
4478+
tool_context,
4479+
dummy_arrow_schema,
4480+
):
4481+
"""End-to-end: after_tool_callback emits tool_origin=TRANSFER_A2A."""
4482+
from google.adk.tools.transfer_to_agent_tool import TransferToAgentTool
4483+
4484+
tool = TransferToAgentTool(
4485+
agent_names=["local", "remote"],
4486+
target_origin_by_name={
4487+
"local": "TRANSFER_AGENT",
4488+
"remote": "TRANSFER_A2A",
4489+
},
4490+
)
4491+
bigquery_agent_analytics_plugin.TraceManager.push_span(tool_context)
4492+
await bq_plugin_inst.after_tool_callback(
4493+
tool=tool,
4494+
tool_args={"agent_name": "remote"},
4495+
tool_context=tool_context,
4496+
result={},
4497+
)
4498+
await asyncio.sleep(0.01)
4499+
log_entry = await _get_captured_event_dict_async(
4500+
mock_write_client, dummy_arrow_schema
4501+
)
4502+
_assert_common_fields(log_entry, "TOOL_COMPLETED")
4503+
content_dict = json.loads(log_entry["content"])
4504+
assert content_dict["tool"] == "transfer_to_agent"
4505+
assert content_dict["tool_origin"] == "TRANSFER_A2A"
4506+
44404507

44414508
class TestHITLTracing:
44424509
"""Tests for HITL-specific event emission via on_event_callback.

0 commit comments

Comments
 (0)