chore: import upstream snapshot with attribution

tests/unittests/workflow/workflow_testing_utils.py

@@ -0,0 +1,399 @@
+# 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.
+
+"""Testing utils for the Workflow."""
+
+import copy
+import inspect
+from typing import Any
+from typing import AsyncGenerator
+from typing import Callable
+from typing import List
+from typing import Optional
+
+from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.context import Context
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.invocation_context import InvocationContext as BaseInvocationContext
+from google.adk.apps.app import ResumabilityConfig
+from google.adk.events.event import Event
+from google.adk.events.request_input import RequestInput
+from google.adk.runners import Runner
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.workflow import BaseNode
+from google.adk.workflow._graph import RouteValue
+from google.adk.workflow.utils._workflow_hitl_utils import has_auth_request_function_call
+from google.adk.workflow.utils._workflow_hitl_utils import has_request_input_function_call
+from google.genai import types
+from pydantic import ConfigDict
+from pydantic import Field
+from typing_extensions import override
+
+from .testing_utils import END_OF_AGENT
+from .testing_utils import simplify_content
+
+
+async def run_workflow(wf, message='start'):
+  """Run a Workflow through Runner, return collected events."""
+  ss = InMemorySessionService()
+  runner = Runner(app_name=wf.name, node=wf, session_service=ss)
+  session = await ss.create_session(app_name=wf.name, user_id='u')
+  msg = types.Content(parts=[types.Part(text=message)], role='user')
+  events = []
+  async for event in runner.run_async(
+      user_id='u', session_id=session.id, new_message=msg
+  ):
+    events.append(event)
+  return events, ss, session
+
+
+# Emulates a node that outputs an Event & a route.
+# If output is not None, the output is set as the data field in the event.
+# If route is not None, the route is set in the node output event.
+# The route can be set without the output. This means didn't produce any output
+# but wants to signal a route to take.
+class TestingNode(BaseNode):
+  model_config = ConfigDict(arbitrary_types_allowed=True)
+
+  output: Optional[Any] = None
+  route: (
+      RouteValue | list[RouteValue] | Callable[[Context, Any], Any] | None
+  ) = None
+  received_inputs: List[Any] = Field(default_factory=list)
+
+  @override
+  async def _run_impl(
+      self,
+      *,
+      ctx: Context,
+      node_input: Any,
+  ) -> AsyncGenerator[Any, None]:
+    if self.output is not None or self.route is not None:
+      route = None
+      if callable(self.route):
+        if inspect.iscoroutinefunction(self.route):
+          route = await self.route(ctx, node_input)
+        else:
+          route = self.route(ctx, node_input)
+      else:
+        route = self.route
+
+      self.received_inputs.append(node_input)
+      yield Event(
+          output=self.output,
+          route=route,
+      )
+
+
+class TestingNodeWithIntermediateContent(BaseNode):
+  model_config = ConfigDict(arbitrary_types_allowed=True)
+
+  intermediate_content: list[types.Content] = Field(default_factory=list)
+  output: Optional[Any] = None
+  route: Optional[str] = None
+
+  @override
+  async def _run_impl(
+      self,
+      *,
+      ctx: Context,
+      node_input: Any,
+  ) -> AsyncGenerator[Any, None]:
+    for content in self.intermediate_content:
+      yield Event(
+          author=self.name,
+          invocation_id=ctx.invocation_id,
+          content=content,
+      )
+
+    if self.output is not None:
+      yield Event(
+          output=self.output,
+          route=self.route,
+      )
+
+
+class InputCapturingNode(BaseNode):
+  """A node that captures the inputs it receives."""
+
+  model_config = ConfigDict(arbitrary_types_allowed=True)
+
+  received_inputs: List[Any] = Field(default_factory=list)
+
+  @override
+  async def _run_impl(
+      self,
+      *,
+      ctx: Context,
+      node_input: Any,
+  ) -> AsyncGenerator[Any, None]:
+    self.received_inputs.append(node_input)
+    yield Event(
+        output={'received': node_input},
+    )
+
+
+class RequestInputNode(BaseNode):
+  """A simple node that requests input from the user."""
+
+  model_config = ConfigDict(arbitrary_types_allowed=True)
+
+  message: str = Field(default='')
+  response_schema: Optional[dict[str, Any]] = None
+
+  @override
+  async def _run_impl(
+      self,
+      *,
+      ctx: Context,
+      node_input: Any,
+  ) -> AsyncGenerator[Any, None]:
+    yield RequestInput(
+        message=self.message,
+        response_schema=self.response_schema,
+    )
+
+
+async def create_parent_invocation_context(
+    test_name: str, agent: BaseAgent, resumable: bool = False
+) -> InvocationContext:
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name='test_app', user_id='test_user'
+  )
+  return InvocationContext(
+      invocation_id=f'{test_name}_invocation_id',
+      agent=agent,
+      session=session,
+      session_service=session_service,
+      resumability_config=ResumabilityConfig(is_resumable=resumable),
+  )
+
+
+def simplify_event_with_node(
+    event: Event,
+    include_state_delta: bool = False,
+) -> Any | None:
+  if isinstance(event, Event):
+    if (
+        'output' not in event.model_fields_set
+        and not (include_state_delta and event.actions.state_delta)
+        and not event.content
+    ):
+      return None
+
+    # If the event has content, return the simplified content.
+    if event.content:
+      return simplify_content(event.content)
+
+    simplified_event = {}
+
+    # Also simplify event.output if it contains Content.
+    # The tests assume that Content found in event data should be simplified
+    # (IDs stripped) just like event.content. This ensures consistent
+    # assertion behavior.
+    output = event.output
+    if isinstance(output, types.Content):
+      output = copy.deepcopy(output)
+      for part in output.parts:
+        if part.function_call and part.function_call.id:
+          part.function_call.id = None
+        if part.function_response and part.function_response.id:
+          part.function_response.id = None
+    simplified_event['output'] = output
+
+    if include_state_delta and event.actions.state_delta:
+      simplified_event['state_delta'] = event.actions.state_delta
+
+    return simplified_event
+  elif event.content:
+    return simplify_content(event.content)
+
+
+def simplify_events_with_node(
+    events: list[Event],
+    *,
+    include_state_delta: bool = False,
+    include_workflow_output: bool = False,
+) -> list[tuple[str, Any]]:
+  results = []
+
+  # Second pass: Simplify events
+  for event in events:
+    # Optionally skip top-level workflow output events (events emitted
+    # by the Workflow in _finalize_workflow). These events have a
+    # top-level node_path (no '/' separator) and carry output data.
+    if (
+        not include_workflow_output
+        and isinstance(event, Event)
+        and event.output is not None
+        and '/' not in (event.node_info.path or '')
+    ):
+      continue
+
+    simplified_event = simplify_event_with_node(event, include_state_delta)
+    if simplified_event:
+      # Map the author to the source node name if it exists.
+      if hasattr(event, 'node_info') and event.node_info.path:
+        author = event.node_info.path
+      else:
+        author = event.author
+      results.append((author, simplified_event))
+  return results
+
+
+def simplify_events_with_node_and_agent_state(
+    events: list[Event],
+    *,
+    include_state_delta: bool = False,
+    include_inputs_and_triggers: bool = False,
+    include_resume_inputs: bool = False,
+    include_workflow_output: bool = False,
+):
+  fields_to_exclude = {'run_id'}
+  if not include_inputs_and_triggers:
+    fields_to_exclude.add('input')
+  if not include_resume_inputs:
+    fields_to_exclude.add('resume_inputs')
+
+  results = []
+
+  for event in events:
+    # Optionally skip top-level workflow output events.
+    if (
+        not include_workflow_output
+        and isinstance(event, Event)
+        and event.output is not None
+        and '/' not in (event.node_info.path or '')
+    ):
+      continue
+    simplified_event = simplify_event_with_node(event, include_state_delta)
+
+    # Map the author to the source node name if it exists.
+    if hasattr(event, 'node_info') and event.node_info.path:
+      author = event.node_info.path
+    else:
+      author = event.author
+
+    if simplified_event:
+      results.append((author, simplified_event))
+    elif event.actions.end_of_agent:
+      results.append((author, END_OF_AGENT))
+    elif event.actions.agent_state is not None:
+      agent_state = event.actions.agent_state
+      nodes = agent_state.get('nodes', {})
+      simplified_nodes = {}
+      for node_name, node_state in nodes.items():
+        simplified_nodes[node_name] = {
+            k: v
+            for k, v in node_state.items()
+            if k not in fields_to_exclude
+            and (k != 'interrupts' or v)  # Exclude empty interrupts
+            and (k != 'resume_inputs' or v)  # Exclude empty resume_inputs
+        }
+      results.append((author, {'nodes': simplified_nodes}))
+  return results
+
+
+def get_request_input_events(events: list[Any]) -> list[Any]:
+  """Returns a list of request input events from the given list of events."""
+  return [e for e in events if has_request_input_function_call(e)]
+
+
+def get_auth_request_events(events: list[Any]) -> list[Any]:
+  """Returns a list of auth credential request events from the given list."""
+  return [e for e in events if has_auth_request_function_call(e)]
+
+
+def get_output_events(events: list[Any], output: Any = None) -> list[Any]:
+  """Returns a list of events that have output populated."""
+  return [
+      e
+      for e in events
+      if isinstance(e, Event)
+      and e.output is not None
+      and (output is None or e.output == output)
+  ]
+
+
+def get_outputs(events: list[Event]) -> list[Any]:
+  """Extracts output values from events, skipping non-output events."""
+  return [
+      e.output for e in events if isinstance(e, Event) and e.output is not None
+  ]
+
+
+def strip_checkpoint_events(
+    simplified_events: list[tuple[str, Any]],
+) -> list[tuple[str, Any]]:
+  """Strips agent_state checkpoint and end_of_agent events.
+
+  In non-resumable mode, the workflow does not emit checkpoint events
+  or end_of_agent events.  Use this to derive the expected simplified
+  output for non-resumable tests from the resumable expected output.
+  """
+  return [
+      (author, data)
+      for author, data in simplified_events
+      if not (isinstance(data, dict) and 'nodes' in data)
+      and data != END_OF_AGENT
+  ]
+
+
+def find_function_call_event(
+    events: list[Any], name: str | None = None
+) -> Any | None:
+  """Finds the first event containing a function call."""
+  for e in events:
+    if hasattr(e, 'content') and e.content and e.content.parts:
+      for part in e.content.parts:
+        if part.function_call:
+          if name is None or part.function_call.name == name:
+            return e
+  return None
+
+
+class CustomRetryableError(Exception):
+  """A custom error meant to be retried."""
+
+
+class CustomNonRetryableError(Exception):
+  """A custom error not meant to be retried."""
+
+
+class _FlakyNode(BaseNode):
+  model_config = ConfigDict(arbitrary_types_allowed=True)
+
+  message: str = Field(default='')
+  succeed_on_iteration: int = Field(default=0)
+  tracker: dict[str, Any] = Field(default_factory=dict)
+  exception_to_raise: Exception = Field(...)
+
+  @override
+  async def run(
+      self,
+      *,
+      ctx: Context,
+      node_input: Any,
+  ) -> AsyncGenerator[Any, None]:
+    iteration_count = self.tracker.get('iteration_count', 0) + 1
+    self.tracker['iteration_count'] = iteration_count
+    self.tracker.setdefault('attempt_counts', []).append(ctx.attempt_count)
+
+    if iteration_count < self.succeed_on_iteration:
+      raise self.exception_to_raise
+
+    yield Event(
+        output=self.message,
+    )