Function test_auto_to_loop

tests/unittests/workflow/test_agent_transfer.py:409–528

(is_resumable: bool)

@pytest.mark.parametrize('is_resumable', [True, False])
def test_auto_to_loop(is_resumable: bool):
  response = [
      transfer_call_part('sub_agent_1'),
      # sub_agent_1 responds directly instead of transferring.
      'response1',
      'response2',
      'response3',
      Part.from_function_call(name='exit_loop', args={}),
      'response4',
      'response5',
  ]
  mock_model = testing_utils.MockModel.create(responses=response)
  # root (auto) - sub_agent_1 (loop) - sub_agent_1_1 (single)
  #                             \ sub_agent_1_2 (single)
  sub_agent_1_1 = Agent(
      name='sub_agent_1_1',
      model=mock_model,
      disallow_transfer_to_parent=True,
      disallow_transfer_to_peers=True,
  )
  sub_agent_1_2 = Agent(
      name='sub_agent_1_2',
      model=mock_model,
      disallow_transfer_to_parent=True,
      disallow_transfer_to_peers=True,
      tools=[exit_loop],
  )
  sub_agent_1 = LoopAgent(
      name='sub_agent_1',
      sub_agents=[sub_agent_1_1, sub_agent_1_2],
  )
  root_agent = Agent(
      name='root_agent',
      model=mock_model,
      sub_agents=[sub_agent_1],
  )
  app = App(
      name='test_app',
      root_agent=root_agent,
      resumability_config=ResumabilityConfig(is_resumable=is_resumable),
  )
  runner = testing_utils.InMemoryRunner(app=app)

  if not is_resumable:
    # Asserts the transfer.
    assert testing_utils.simplify_events(runner.run('test1')) == [
        # Transfers to sub_agent_1.
        ('root_agent', transfer_call_part('sub_agent_1')),
        ('root_agent', TRANSFER_RESPONSE_PART),
        # Loops.
        ('sub_agent_1_1', 'response1'),
        ('sub_agent_1_2', 'response2'),
        ('sub_agent_1_1', 'response3'),
        # Exits.
        ('sub_agent_1_2', Part.from_function_call(name='exit_loop', args={})),
        (
            'sub_agent_1_2',
            Part.from_function_response(