test: Make testing_utils.InMemoryRunner support ADK App and add utils for extracting event contents for testing resumability

PiperOrigin-RevId: 811933527

tests/unittests/testing_utils.py

@@ -16,6 +16,7 @@ import asyncio
 import contextlib
 from typing import AsyncGenerator
 from typing import Generator
+from typing import Optional
 from typing import Union
 
 from google.adk.agents.invocation_context import InvocationContext
@@ -23,6 +24,7 @@ from google.adk.agents.live_request_queue import LiveRequestQueue
 from google.adk.agents.llm_agent import Agent
 from google.adk.agents.llm_agent import LlmAgent
 from google.adk.agents.run_config import RunConfig
+from google.adk.apps.app import App
 from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
 from google.adk.events.event import Event
 from google.adk.memory.in_memory_memory_service import InMemoryMemoryService
@@ -119,7 +121,32 @@ def append_user_content(
 # Extracts the contents from the events and transform them into a list of
 # (author, simplified_content) tuples.
 def simplify_events(events: list[Event]) -> list[(str, types.Part)]:
-  return [(event.author, simplify_content(event.content)) for event in events]
+  return [
+      (event.author, simplify_content(event.content))
+      for event in events
+      if event.content
+  ]
+
+
+END_OF_AGENT = 'end_of_agent'
+
+
+# Extracts the contents from the events and transform them into a list of
+# (author, simplified_content OR AgentState OR "end_of_agent") tuples.
+#
+# Could be used to compare events for testing resumability.
+def simplify_resumable_app_events(
+    events: list[Event],
+) -> list[(str, Union[types.Part, str])]:
+  results = []
+  for event in events:
+    if event.content:
+      results.append((event.author, simplify_content(event.content)))
+    elif event.actions.end_of_agent:
+      results.append((event.author, END_OF_AGENT))
+    elif event.actions.agent_state is not None:
+      results.append((event.author, event.actions.agent_state))
+  return results
 
 
 # Simplifies the contents into a list of (author, simplified_content) tuples.
@@ -189,31 +216,52 @@ class InMemoryRunner:
 
   def __init__(
       self,
-      root_agent: Union[Agent, LlmAgent],
+      root_agent: Optional[Union[Agent, LlmAgent]] = None,
       response_modalities: list[str] = None,
       plugins: list[BasePlugin] = [],
+      app: Optional[App] = None,
   ):
-    self.root_agent = root_agent
-    self.runner = Runner(
-        app_name='test_app',
-        agent=root_agent,
-        artifact_service=InMemoryArtifactService(),
-        session_service=InMemorySessionService(),
-        memory_service=InMemoryMemoryService(),
-        plugins=plugins,
-    )
+    """Initializes the InMemoryRunner.
+
+    Args:
+      root_agent: The root agent to run, won't be used if app is provided.
+      response_modalities: The response modalities of the runner.
+      plugins: The plugins to use in the runner, won't be used if app is
+        provided.
+      app: The app to use in the runner.
+    """
+    if not app:
+      self.app_name = 'test_app'
+      self.root_agent = root_agent
+      self.runner = Runner(
+          app_name='test_app',
+          agent=root_agent,
+          artifact_service=InMemoryArtifactService(),
+          session_service=InMemorySessionService(),
+          memory_service=InMemoryMemoryService(),
+          plugins=plugins,
+      )
+    else:
+      self.app_name = app.name
+      self.root_agent = app.root_agent
+      self.runner = Runner(
+          app=app,
+          artifact_service=InMemoryArtifactService(),
+          session_service=InMemorySessionService(),
+          memory_service=InMemoryMemoryService(),
+      )
     self.session_id = None
 
   @property
   def session(self) -> Session:
     if not self.session_id:
       session = self.runner.session_service.create_session_sync(
-          app_name='test_app', user_id='test_user'
+          app_name=self.app_name, user_id='test_user'
       )
       self.session_id = session.id
       return session
     return self.runner.session_service.get_session_sync(
-        app_name='test_app', user_id='test_user', session_id=self.session_id
+        app_name=self.app_name, user_id='test_user', session_id=self.session_id
     )
 
   def run(self, new_message: types.ContentUnion) -> list[Event]: