Function test_run_async

tests/unittests/agents/test_parallel_agent.py:78–121

(request: pytest.FixtureRequest, is_resumable: bool)

@pytest.mark.asyncio
@pytest.mark.parametrize('is_resumable', [True, False])
async def test_run_async(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')
  parallel_agent = ParallelAgent(
      name=f'{request.function.__name__}_test_parallel_agent',
      sub_agents=[
          agent1,
          agent2,
      ],
  )
  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) == 4

    assert events[0].author == parallel_agent.name
    assert not events[0].actions.end_of_agent

    # agent2 generates an event first, then agent1. Because they run in parallel
    # and agent1 has a delay.
    assert events[1].author == agent2.name
    assert events[2].author == agent1.name
    assert events[1].branch == f'{parallel_agent.name}.{agent2.name}'
    assert events[2].branch == f'{parallel_agent.name}.{agent1.name}'
    assert events[1].content.parts[0].text == f'Hello, async {agent2.name}!'
    assert events[2].content.parts[0].text == f'Hello, async {agent1.name}!'

    assert events[3].author == parallel_agent.name
    assert events[3].actions.end_of_agent
  else:
    assert len(events) == 2

    assert events[0].author == agent2.name
    assert events[1].author == agent1.name
    assert events[0].branch == f'{parallel_agent.name}.{agent2.name}'
    assert events[1].branch == f'{parallel_agent.name}.{agent1.name}'
    assert events[0].content.parts[0].text == f'Hello, async {agent2.name}!'
    assert events[1].content.parts[0].text == f'Hello, async {agent1.name}!'


@pytest.mark.asyncio