Function test_run_async_branches

tests/unittests/agents/test_parallel_agent.py:126–212

(
    request: pytest.FixtureRequest, is_resumable: bool
)

@pytest.mark.asyncio
@pytest.mark.parametrize('is_resumable', [True, False])
async def test_run_async_branches(
    request: pytest.FixtureRequest, is_resumable: bool
):
  agent1 = _TestingAgent(
      name=f'{request.function.__name__}_test_agent_1',
      delay=0.5,
  )
  agent2 = _TestingAgent(name=f'{request.function.__name__}_test_agent_2')
  agent3 = _TestingAgent(name=f'{request.function.__name__}_test_agent_3')
  sequential_agent = SequentialAgent(
      name=f'{request.function.__name__}_test_sequential_agent',
      sub_agents=[agent2, agent3],
  )
  parallel_agent = ParallelAgent(
      name=f'{request.function.__name__}_test_parallel_agent',
      sub_agents=[
          sequential_agent,
          agent1,
      ],
  )
  parent_ctx = await _create_parent_invocation_context(
      request.function.__name__, parallel_agent, is_resumable=is_resumable
  )
  events = [e async for e in parallel_agent.run_async(parent_ctx)]

  if is_resumable:
    assert len(events) == 8

    # 1. parallel agent checkpoint
    assert events[0].author == parallel_agent.name
    assert not events[0].actions.end_of_agent

    # 2. sequential agent checkpoint
    assert events[1].author == sequential_agent.name
    assert not events[1].actions.end_of_agent
    assert events[1].actions.agent_state['current_sub_agent'] == agent2.name
    assert events[1].branch == f'{parallel_agent.name}.{sequential_agent.name}'

    # 3. agent 2 event
    assert events[2].author == agent2.name
    assert events[2].branch == f'{parallel_agent.name}.{sequential_agent.name}'

    # 4. sequential agent checkpoint
    assert events[3].author == sequential_agent.name
    assert not events[3].actions.end_of_agent
    assert events[3].actions.agent_state['current_sub_agent'] == agent3.name
    assert events[3].branch == f'{parallel_agent.name}.{sequential_agent.name}'

    # 5. agent 3 event
    assert events[4].author == agent3.name
    assert events[4].branch == f'{parallel_agent.name}.{sequential_agent.name}'

    # 6. sequential agent checkpoint (end)
    assert events[5].author == sequential_agent.name
    assert events[5].actions.end_of_agent
    assert events[5].branch == f'{parallel_agent.name}.{sequential_agent.name}'

    # Descendants of the same sub-agent should have the same branch.