chore: add test for parallel agent to verify correct ordering of agents

PiperOrigin-RevId: 796949367

Author: jawoszek

tests/unittests/agents/test_parallel_agent.py

@@ -135,3 +135,63 @@ async def test_run_async_branches(request: pytest.FixtureRequest):
   # Sub-agents should have different branches.
   assert events[2].branch != events[1].branch
   assert events[2].branch != events[0].branch
+
+
+class _TestingAgentWithMultipleEvents(BaseAgent):
+  """Mock agent for testing."""
+
+  processed_events: int = 0
+
+  @override
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    event = Event(
+        author=self.name,
+        branch=ctx.branch,
+        invocation_id=ctx.invocation_id,
+        content=types.Content(
+            parts=[types.Part(text=f'Hello, async {self.name}!')]
+        ),
+    )
+    for _ in range(0, 3):
+      yield event
+      self.processed_events += 1
+
+
+@pytest.mark.asyncio
+async def test_generating_one_event_per_agent_at_once(
+    request: pytest.FixtureRequest,
+):
+  # This test is to verify that the parallel agent won't generate more than one
+  # event per agent at a time.
+  agent1 = _TestingAgentWithMultipleEvents(
+      name=f'{request.function.__name__}_test_agent_1'
+  )
+  agent2 = _TestingAgentWithMultipleEvents(
+      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
+  )
+
+  agen = parallel_agent.run_async(parent_ctx)
+  # No event is generated yet.
+  assert agent1.processed_events + agent2.processed_events == 0
+  current_iter = 0
+  async for _ in agen:
+    processed_events = agent1.processed_events + agent2.processed_events
+    # Depending on coroutine scheduling, agent increments the counter before
+    # or after next iteration, so we have two possible/expected values of sum.
+    assert (
+        processed_events == current_iter or processed_events == current_iter - 1
+    )
+    # No more than one event is generated per agent after each iteration.
+    current_iter += 1