refactor: Extract reusable private methods

Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com>
PiperOrigin-RevId: 871128393
This commit is contained in:
Xiang (Sean) Zhou
2026-02-16 22:52:47 -08:00
committed by Copybara-Service
parent 4e2d6159ae
commit 706f9fe74d
4 changed files with 384 additions and 334 deletions
@@ -19,6 +19,7 @@ from unittest.mock import AsyncMock
from google.adk.agents.llm_agent import Agent
from google.adk.events.event import Event
from google.adk.flows.llm_flows.base_llm_flow import _handle_after_model_callback
from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow
from google.adk.models.google_llm import Gemini
from google.adk.models.llm_request import LlmRequest
@@ -285,7 +286,7 @@ async def test_handle_after_model_callback_grounding_with_no_callbacks(
)
flow = BaseLlmFlowForTesting()
result = await flow._handle_after_model_callback(
result = await _handle_after_model_callback(
invocation_context, llm_response, event
)
@@ -342,7 +343,7 @@ async def test_handle_after_model_callback_grounding_with_callback_override(
)
flow = BaseLlmFlowForTesting()
result = await flow._handle_after_model_callback(
result = await _handle_after_model_callback(
invocation_context, llm_response, event
)
@@ -404,7 +405,7 @@ async def test_handle_after_model_callback_grounding_with_plugin_override(
)
flow = BaseLlmFlowForTesting()
result = await flow._handle_after_model_callback(
result = await _handle_after_model_callback(
invocation_context, llm_response, event
)
@@ -461,13 +462,13 @@ async def test_handle_after_model_callback_caches_canonical_tools():
flow = BaseLlmFlowForTesting()
# Call _handle_after_model_callback multiple times with the same context
result1 = await flow._handle_after_model_callback(
result1 = await _handle_after_model_callback(
invocation_context, llm_response, event
)
result2 = await flow._handle_after_model_callback(
result2 = await _handle_after_model_callback(
invocation_context, llm_response, event
)
result3 = await flow._handle_after_model_callback(
result3 = await _handle_after_model_callback(
invocation_context, llm_response, event
)