chore: import upstream snapshot with attribution

python/tests/unit/agents/orchestration/test_concurrent.py

@@ -0,0 +1,204 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import sys
+from unittest.mock import patch
+
+import pytest
+
+from semantic_kernel.agents.orchestration.concurrent import ConcurrentOrchestration
+from semantic_kernel.agents.orchestration.orchestration_base import DefaultTypeAlias, OrchestrationResult
+from semantic_kernel.agents.runtime.in_process.in_process_runtime import InProcessRuntime
+from semantic_kernel.contents.chat_message_content import ChatMessageContent
+from semantic_kernel.contents.streaming_chat_message_content import StreamingChatMessageContent
+from tests.unit.agents.orchestration.conftest import MockAgent, MockAgentWithException, MockRuntime
+
+
+async def test_prepare():
+    """Test the prepare method of the ConcurrentOrchestration."""
+    agent_a = MockAgent()
+    agent_b = MockAgent()
+
+    runtime = MockRuntime()
+
+    package_path = "semantic_kernel.agents.orchestration.concurrent"
+    with (
+        patch(f"{package_path}.ConcurrentOrchestration._start"),
+        patch(f"{package_path}.ConcurrentAgentActor.register") as mock_agent_actor_register,
+        patch(f"{package_path}.CollectionActor.register") as mock_collection_actor_register,
+        patch.object(runtime, "add_subscription") as mock_add_subscription,
+    ):
+        orchestration = ConcurrentOrchestration(members=[agent_a, agent_b])
+        await orchestration.invoke(task="test_message", runtime=runtime)
+
+        assert mock_agent_actor_register.call_count == 2
+        assert mock_collection_actor_register.call_count == 1
+        assert mock_add_subscription.call_count == 2
+
+
+async def test_invoke():
+    """Test the invoke method of the ConcurrentOrchestration."""
+    agent_a = MockAgent()
+    agent_b = MockAgent()
+
+    runtime = InProcessRuntime()
+    runtime.start()
+
+    try:
+        orchestration = ConcurrentOrchestration(members=[agent_a, agent_b])
+        orchestration_result = await orchestration.invoke(task="test_message", runtime=runtime)
+        result = await orchestration_result.get(1.0)
+
+        assert isinstance(orchestration_result, OrchestrationResult)
+        assert isinstance(result, list)
+        assert len(result) == 2
+        assert all(isinstance(item, ChatMessageContent) for item in result)
+    finally:
+        await runtime.stop_when_idle()
+
+
+async def test_invoke_with_response_callback():
+    """Test the invoke method of the ConcurrentOrchestration with a response callback."""
+    agent_a = MockAgent()
+    agent_b = MockAgent()
+
+    runtime = InProcessRuntime()
+    runtime.start()
+
+    responses: list[DefaultTypeAlias] = []
+    try:
+        orchestration = ConcurrentOrchestration(
+            members=[agent_a, agent_b],
+            agent_response_callback=lambda x: responses.append(x),
+        )
+        orchestration_result = await orchestration.invoke(task="test_message", runtime=runtime)
+        await orchestration_result.get(1.0)
+
+        assert len(responses) == 2
+        assert all(isinstance(item, ChatMessageContent) for item in responses)
+        assert all(item.content == "mock_response" for item in responses)
+    finally:
+        await runtime.stop_when_idle()
+
+
+async def test_invoke_with_streaming_response_callback():
+    """Test the invoke method of the ConcurrentOrchestration with a streaming response callback."""
+    agent_a = MockAgent()
+    agent_b = MockAgent()
+
+    runtime = InProcessRuntime()
+    runtime.start()
+
+    responses: dict[str, list[StreamingChatMessageContent]] = {}
+    try:
+        orchestration = ConcurrentOrchestration(
+            members=[agent_a, agent_b],
+            streaming_agent_response_callback=lambda x, _: responses.setdefault(x.name, []).append(x),
+        )
+        orchestration_result = await orchestration.invoke(task="test_message", runtime=runtime)
+        await orchestration_result.get(1.0)
+
+        assert len(responses[agent_a.name]) == 2
+        assert len(responses[agent_b.name]) == 2
+
+        assert all(isinstance(item, StreamingChatMessageContent) for item in responses[agent_a.name])
+        assert all(isinstance(item, StreamingChatMessageContent) for item in responses[agent_b.name])
+
+        agent_a_response = sum(responses[agent_a.name][1:], responses[agent_a.name][0])
+        assert agent_a_response.content == "mock_response"
+        agent_b_response = sum(responses[agent_b.name][1:], responses[agent_b.name][0])
+        assert agent_b_response.content == "mock_response"
+    finally:
+        await runtime.stop_when_idle()
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 11),
+    reason="Python 3.10 doesn't bound the original function provided to the wraps argument of the patch object.",
+)
+async def test_invoke_cancel_before_completion():
+    """Test the invoke method of the ConcurrentOrchestration with cancellation before completion."""
+    with (
+        patch.object(MockAgent, "invoke_stream", wraps=MockAgent.invoke_stream, autospec=True) as mock_invoke_stream,
+    ):
+        agent_a = MockAgent()
+        agent_b = MockAgent()
+
+        runtime = InProcessRuntime()
+        runtime.start()
+
+        try:
+            orchestration = ConcurrentOrchestration(members=[agent_a, agent_b])
+            orchestration_result = await orchestration.invoke(task="test_message", runtime=runtime)
+
+            # Cancel before the collection agent gets the responses
+            await asyncio.sleep(0.05)
+            orchestration_result.cancel()
+        finally:
+            await runtime.stop_when_idle()
+
+        assert mock_invoke_stream.call_count == 2
+
+
+async def test_invoke_cancel_after_completion():
+    """Test the invoke method of the ConcurrentOrchestration with cancellation after completion."""
+    agent_a = MockAgent()
+    agent_b = MockAgent()
+
+    runtime = InProcessRuntime()
+    runtime.start()
+
+    try:
+        orchestration = ConcurrentOrchestration(members=[agent_a, agent_b])
+        orchestration_result = await orchestration.invoke(task="test_message", runtime=runtime)
+
+        # Wait for the orchestration to complete
+        await orchestration_result.get(1.0)
+
+        with pytest.raises(RuntimeError, match="The invocation has already been completed."):
+            orchestration_result.cancel()
+    finally:
+        await runtime.stop_when_idle()
+
+
+async def test_invoke_with_double_get_result():
+    """Test the invoke method of the ConcurrentOrchestration with double get result."""
+    agent_a = MockAgent()
+    agent_b = MockAgent()
+
+    runtime = InProcessRuntime()
+    runtime.start()
+
+    try:
+        orchestration = ConcurrentOrchestration(members=[agent_a, agent_b])
+        orchestration_result = await orchestration.invoke(task="test_message", runtime=runtime)
+
+        # Get result before completion
+        with pytest.raises(asyncio.TimeoutError):
+            await orchestration_result.get(0.1)
+        # The invocation should still be in progress and getting the result again should not raise an error
+        result = await orchestration_result.get(1.0)
+
+        assert isinstance(result, list)
+        assert len(result) == 2
+    finally:
+        await runtime.stop_when_idle()
+
+
+async def test_invoke_with_agent_raising_exception():
+    """Test the invoke method of the ConcurrentOrchestration with an agent raising an exception."""
+    agent_a = MockAgent()
+    agent_b = MockAgentWithException()
+
+    runtime = InProcessRuntime()
+    runtime.start()
+
+    try:
+        orchestration = ConcurrentOrchestration(members=[agent_a, agent_b])
+        orchestration_result = await orchestration.invoke(task="test_message", runtime=runtime)
+
+        with pytest.raises(RuntimeError, match="Mock agent exception"):
+            await orchestration_result.get(1.0)
+        assert orchestration_result.exception is not None
+    finally:
+        await runtime.stop_when_idle()