refactor: Extract reusable private methods

Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com>
PiperOrigin-RevId: 871128393

tests/unittests/auth/test_toolset_auth.py

@@ -31,6 +31,7 @@ from google.adk.auth.auth_credential import OAuth2Auth
 from google.adk.auth.auth_preprocessor import TOOLSET_AUTH_CREDENTIAL_ID_PREFIX
 from google.adk.auth.auth_tool import AuthConfig
 from google.adk.auth.auth_tool import AuthToolArguments
+from google.adk.flows.llm_flows.base_llm_flow import _resolve_toolset_auth
 from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow
 from google.adk.flows.llm_flows.base_llm_flow import TOOLSET_AUTH_CREDENTIAL_ID_PREFIX as FLOW_PREFIX
 from google.adk.flows.llm_flows.functions import build_auth_request_event
@@ -119,14 +120,6 @@ class TestResolveToolsetAuth:
     agent.tools = []
     return agent
 
-  @pytest.fixture
-  def flow(self):
-    """Create a BaseLlmFlow instance for testing."""
-    # BaseLlmFlow is abstract, but we can still test _resolve_toolset_auth
-    flow = Mock(spec=BaseLlmFlow)
-    flow._resolve_toolset_auth = BaseLlmFlow._resolve_toolset_auth
-    return flow
-
   @pytest.mark.asyncio
   async def test_no_tools_returns_no_events(
       self, mock_invocation_context, mock_agent
@@ -134,9 +127,8 @@ class TestResolveToolsetAuth:
     """Test that no events are yielded when agent has no tools."""
     mock_agent.tools = []
 
-    flow = BaseLlmFlow.__new__(BaseLlmFlow)
     events = []
-    async for event in flow._resolve_toolset_auth(
+    async for event in _resolve_toolset_auth(
         mock_invocation_context, mock_agent
     ):
       events.append(event)
@@ -152,9 +144,8 @@ class TestResolveToolsetAuth:
     toolset = MockToolset(auth_config=None)
     mock_agent.tools = [toolset]
 
-    flow = BaseLlmFlow.__new__(BaseLlmFlow)
     events = []
-    async for event in flow._resolve_toolset_auth(
+    async for event in _resolve_toolset_auth(
         mock_invocation_context, mock_agent
     ):
       events.append(event)
@@ -184,9 +175,8 @@ class TestResolveToolsetAuth:
       mock_manager.get_auth_credential = AsyncMock(return_value=mock_credential)
       MockCredentialManager.return_value = mock_manager
 
-      flow = BaseLlmFlow.__new__(BaseLlmFlow)
       events = []
-      async for event in flow._resolve_toolset_auth(
+      async for event in _resolve_toolset_auth(
           mock_invocation_context, mock_agent
       ):
         events.append(event)
@@ -213,9 +203,8 @@ class TestResolveToolsetAuth:
       mock_manager.get_auth_credential = AsyncMock(return_value=None)
       MockCredentialManager.return_value = mock_manager
 
-      flow = BaseLlmFlow.__new__(BaseLlmFlow)
       events = []
-      async for event in flow._resolve_toolset_auth(
+      async for event in _resolve_toolset_auth(
           mock_invocation_context, mock_agent
       ):
         events.append(event)
@@ -258,9 +247,8 @@ class TestResolveToolsetAuth:
       mock_manager.get_auth_credential = AsyncMock(return_value=None)
       MockCredentialManager.return_value = mock_manager
 
-      flow = BaseLlmFlow.__new__(BaseLlmFlow)
       events = []
-      async for event in flow._resolve_toolset_auth(
+      async for event in _resolve_toolset_auth(
           mock_invocation_context, mock_agent
       ):
         events.append(event)

tests/unittests/flows/llm_flows/test_base_llm_flow.py

@@ -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
     )