test(workflow): run the invariant plugin in test modules that build a raw Runner

`InvariantPlugin` checks event/session invariants on every run, but it
only runs when a test goes through the shared test runners. Several
modules build a bare `Runner(...)` because they own the session service
or pass an `App`, so their runs were never checked.

Add `workflow_testing_utils.create_runner(**runner_kwargs,
check_invariants=True)`, which builds a `Runner` with the plugin
installed first (mirroring `TestInMemoryRunner`), and route the raw
sites through it:

- `test_workflow_dynamic_nodes.py` (24), `test_workflow_live.py` (6),
  `test_workflow_rewind.py` (3), `test_workflow_agent_as_node.py` (2),
  `test_workflow_function_tool_as_node.py` (1),
  `test_node_runner_failure.py` (9, plus its local `_run_workflow` copy
  is replaced by the shared `run_workflow`).
- `agents/test_llm_agent_interruptions.py` now uses
  `TestInMemoryRunner` via a small `_make_runner` helper (4 sites).

46 previously unchecked runner sites now run under the plugin; all
migrated modules pass with no new invariant violations, so no opt-outs
were needed. Test-only change.

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 995957207

Author: DeanChensj

tests/unittests/agents/test_llm_agent_interruptions.py

@@ -15,8 +15,6 @@
 from __future__ import annotations
 
 from google.adk.agents import LlmAgent
-from google.adk.runners import Runner
-from google.adk.sessions.in_memory_session_service import InMemorySessionService
 from google.adk.tools.function_tool import FunctionTool
 from google.adk.tools.long_running_tool import LongRunningFunctionTool
 from google.adk.tools.tool_context import ToolContext
@@ -30,6 +28,15 @@
 _SESSION_ID = 'test_session'
 
 
+async def _make_runner(agent):
+  """Build an invariant-checking runner with a pre-created session."""
+  runner = testing_utils.TestInMemoryRunner(agent=agent, app_name='test')
+  await runner.session_service.create_session(
+      app_name='test', user_id=_USER_ID, session_id=_SESSION_ID
+  )
+  return runner
+
+
 async def _setup_runner(mock_model, tools=None, **agent_kwargs):
   """Setup runner with LlmAgent directly."""
   llm_agent = LlmAgent(
@@ -38,16 +45,7 @@ async def _setup_runner(mock_model, tools=None, **agent_kwargs):
       tools=tools or [],
       **agent_kwargs,
   )
-  session_service = InMemorySessionService()
-  await session_service.create_session(
-      app_name='test', user_id=_USER_ID, session_id=_SESSION_ID
-  )
-  runner = Runner(
-      app_name='test',
-      agent=llm_agent,
-      session_service=session_service,
-  )
-  return runner
+  return await _make_runner(llm_agent)
 
 
 async def _run_turn(runner, user_message):
@@ -332,14 +330,7 @@ def transfer_to_child(tool_context: ToolContext) -> str:
         sub_agents=[child_agent],
     )
 
-    # Setup runner
-    session_service = InMemorySessionService()
-    await session_service.create_session(
-        app_name='test', user_id=_USER_ID, session_id=_SESSION_ID
-    )
-    runner = Runner(
-        app_name='test', agent=parent_agent, session_service=session_service
-    )
+    runner = await _make_runner(parent_agent)
 
     # When Parent runs the first turn
     events = await _run_turn(runner, 'Go')
@@ -398,14 +389,7 @@ def transfer_to_child(tool_context: ToolContext) -> str:
         sub_agents=[child_agent],
     )
 
-    # Setup runner
-    session_service = InMemorySessionService()
-    await session_service.create_session(
-        app_name='test', user_id=_USER_ID, session_id=_SESSION_ID
-    )
-    runner = Runner(
-        app_name='test', agent=parent_agent, session_service=session_service
-    )
+    runner = await _make_runner(parent_agent)
 
     # When Parent runs the first turn
     events = await _run_turn(runner, 'Go')
@@ -459,14 +443,7 @@ async def test_single_turn_child_agent_interrupt_and_resume(self):
         sub_agents=[child_agent],
     )
 
-    # Setup runner
-    session_service = InMemorySessionService()
-    await session_service.create_session(
-        app_name='test', user_id=_USER_ID, session_id=_SESSION_ID
-    )
-    runner = Runner(
-        app_name='test', agent=parent_agent, session_service=session_service
-    )
+    runner = await _make_runner(parent_agent)
 
     # When Parent runs the first turn
     events = await _run_turn(runner, 'Go')

tests/unittests/workflow/test_node_runner_failure.py

@@ -24,7 +24,6 @@
 from google.adk.agents.context import Context
 from google.adk.agents.invocation_context import InvocationContext
 from google.adk.events.event import Event
-from google.adk.runners import Runner
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
 from google.adk.sessions.session import Session
 from google.adk.workflow import BaseNode
@@ -45,30 +44,15 @@
 import pytest
 from typing_extensions import override
 
-from .._invariants import InvariantPlugin
 from .workflow_testing_utils import _FlakyNode
+from .workflow_testing_utils import create_runner
 from .workflow_testing_utils import CustomNonRetryableError
 from .workflow_testing_utils import CustomRetryableError
+from .workflow_testing_utils import run_workflow as _run_workflow
 from .workflow_testing_utils import simplify_events_with_node
 from .workflow_testing_utils import TestingNode
 
 
-async def _run_workflow(wf, message='start', *, check_invariants: bool = True):
-  """Run a Workflow through Runner, return collected events."""
-  ss = InMemorySessionService()
-  runner = Runner(app_name=wf.name, node=wf, session_service=ss)
-  if check_invariants:
-    runner.plugin_manager.plugins.insert(0, InvariantPlugin())
-  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
-
-
 @pytest.mark.asyncio
 async def test_node_retries_on_matched_exception_string(
     request: pytest.FixtureRequest,
@@ -170,7 +154,7 @@ async def test_node_fails_immediately_on_unmatched_exception_string(
   )
 
   ss = InMemorySessionService()
-  runner = Runner(app_name=agent.name, node=agent, session_service=ss)
+  runner = create_runner(app_name=agent.name, node=agent, session_service=ss)
   session = await ss.create_session(app_name=agent.name, user_id='u')
   msg = types.Content(parts=[types.Part(text='start')], role='user')
   events = []
@@ -364,7 +348,7 @@ async def test_node_stops_retrying_after_max_attempts(
   )
 
   ss = InMemorySessionService()
-  runner = Runner(app_name=agent.name, node=agent, session_service=ss)
+  runner = create_runner(app_name=agent.name, node=agent, session_service=ss)
   session = await ss.create_session(app_name=agent.name, user_id='u')
   msg = types.Content(parts=[types.Part(text='start')], role='user')
   events = []
@@ -435,7 +419,7 @@ async def test_node_fails_immediately_without_retry_config(
   )
 
   ss = InMemorySessionService()
-  runner = Runner(app_name=agent.name, node=agent, session_service=ss)
+  runner = create_runner(app_name=agent.name, node=agent, session_service=ss)
   session = await ss.create_session(app_name=agent.name, user_id='u')
   msg = types.Content(parts=[types.Part(text='start')], role='user')
   events = []
@@ -563,7 +547,7 @@ async def test_node_waits_for_initial_delay_before_retry(
   )
 
   ss = InMemorySessionService()
-  runner = Runner(app_name=agent.name, node=agent, session_service=ss)
+  runner = create_runner(app_name=agent.name, node=agent, session_service=ss)
   session = await ss.create_session(app_name=agent.name, user_id='u')
   msg = types.Content(parts=[types.Part(text='start')], role='user')
 
@@ -639,7 +623,7 @@ async def test_retry_applies_backoff_strategy(request: pytest.FixtureRequest):
   )
 
   ss = InMemorySessionService()
-  runner = Runner(app_name=agent.name, node=agent, session_service=ss)
+  runner = create_runner(app_name=agent.name, node=agent, session_service=ss)
   session = await ss.create_session(app_name=agent.name, user_id='u')
   msg = types.Content(parts=[types.Part(text='start')], role='user')
 
@@ -718,7 +702,7 @@ async def test_retry_applies_random_jitter(request: pytest.FixtureRequest):
   )
 
   ss = InMemorySessionService()
-  runner = Runner(app_name=agent.name, node=agent, session_service=ss)
+  runner = create_runner(app_name=agent.name, node=agent, session_service=ss)
   session = await ss.create_session(app_name=agent.name, user_id='u')
   msg = types.Content(parts=[types.Part(text='start')], role='user')
 
@@ -926,7 +910,7 @@ async def test_node_fails_immediately_on_unmatched_exception_class(
   )
 
   ss = InMemorySessionService()
-  runner = Runner(app_name=agent.name, node=agent, session_service=ss)
+  runner = create_runner(app_name=agent.name, node=agent, session_service=ss)
   session = await ss.create_session(app_name=agent.name, user_id='u')
   msg = types.Content(parts=[types.Part(text='start')], role='user')
   events = []
@@ -995,7 +979,7 @@ async def test_error_event_emitted_on_failure(
   )
 
   ss = InMemorySessionService()
-  runner = Runner(app_name=agent.name, node=agent, session_service=ss)
+  runner = create_runner(app_name=agent.name, node=agent, session_service=ss)
   session = await ss.create_session(app_name=agent.name, user_id='u')
   msg = types.Content(parts=[types.Part(text='start')], role='user')
   events = []
@@ -1047,7 +1031,7 @@ async def test_error_event_emitted_on_each_retry(
   )
 
   ss = InMemorySessionService()
-  runner = Runner(app_name=agent.name, node=agent, session_service=ss)
+  runner = create_runner(app_name=agent.name, node=agent, session_service=ss)
   session = await ss.create_session(app_name=agent.name, user_id='u')
   msg = types.Content(parts=[types.Part(text='start')], role='user')
   events = []

tests/unittests/workflow/test_workflow_agent_as_node.py

@@ -20,14 +20,14 @@
 from google.adk.agents.context import Context
 from google.adk.agents.invocation_context import InvocationContext as BaseInvocationContext
 from google.adk.events.event import Event
-from google.adk.runners import Runner
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
 from google.adk.sessions.session import Session
 from google.adk.workflow import START
 from google.adk.workflow._workflow import Workflow
 from google.genai import types
 import pytest
 
+from .workflow_testing_utils import create_runner
 from .workflow_testing_utils import InputCapturingNode
 from .workflow_testing_utils import simplify_events_with_node
 
@@ -97,7 +97,7 @@ async def test_run_async_with_agent_nodes(request: pytest.FixtureRequest):
       ],
   )
   ss = InMemorySessionService()
-  runner = Runner(app_name='test', node=wf, session_service=ss)
+  runner = create_runner(app_name='test', node=wf, session_service=ss)
   session = await ss.create_session(app_name='test', user_id='u')
 
   msg = types.Content(parts=[types.Part(text='start')], role='user')
@@ -128,7 +128,7 @@ async def test_run_async_with_agent_node_piping_data(
       ],
   )
   ss = InMemorySessionService()
-  runner = Runner(app_name='test', node=wf, session_service=ss)
+  runner = create_runner(app_name='test', node=wf, session_service=ss)
   session = await ss.create_session(app_name='test', user_id='u')
 
   msg = types.Content(parts=[types.Part(text='start')], role='user')