chore: import upstream snapshot with attribution

python/packages/azurefunctions/tests/test_workflow.py

@@ -0,0 +1,328 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for workflow orchestration functions."""
+
+import json
+from dataclasses import dataclass
+from typing import Any
+
+from agent_framework import (
+    AgentExecutorRequest,
+    AgentExecutorResponse,
+    AgentResponse,
+    Message,
+)
+from agent_framework._workflows._edge import (
+    FanInEdgeGroup,
+    FanOutEdgeGroup,
+    SingleEdgeGroup,
+    SwitchCaseEdgeGroup,
+    SwitchCaseEdgeGroupCase,
+    SwitchCaseEdgeGroupDefault,
+)
+
+from agent_framework_azurefunctions._workflow import (
+    _extract_message_content,
+    build_agent_executor_response,
+    route_message_through_edge_groups,
+)
+
+
+class TestRouteMessageThroughEdgeGroups:
+    """Test suite for route_message_through_edge_groups function."""
+
+    def test_single_edge_group_routes_when_condition_matches(self) -> None:
+        """Test SingleEdgeGroup routes when condition is satisfied."""
+        group = SingleEdgeGroup(source_id="src", target_id="tgt", condition=lambda m: True)
+
+        targets = route_message_through_edge_groups([group], "src", "any message")
+
+        assert targets == ["tgt"]
+
+    def test_single_edge_group_does_not_route_when_condition_fails(self) -> None:
+        """Test SingleEdgeGroup does not route when condition fails."""
+        group = SingleEdgeGroup(source_id="src", target_id="tgt", condition=lambda m: False)
+
+        targets = route_message_through_edge_groups([group], "src", "any message")
+
+        assert targets == []
+
+    def test_single_edge_group_ignores_different_source(self) -> None:
+        """Test SingleEdgeGroup ignores messages from different sources."""
+        group = SingleEdgeGroup(source_id="src", target_id="tgt", condition=lambda m: True)
+
+        targets = route_message_through_edge_groups([group], "other_src", "any message")
+
+        assert targets == []
+
+    def test_switch_case_with_selection_func(self) -> None:
+        """Test SwitchCaseEdgeGroup uses selection_func."""
+
+        def select_first_target(msg: Any, targets: list[str]) -> list[str]:
+            return [targets[0]]
+
+        group = SwitchCaseEdgeGroup(
+            source_id="src",
+            cases=[
+                SwitchCaseEdgeGroupCase(condition=lambda m: True, target_id="target_a"),
+                SwitchCaseEdgeGroupDefault(target_id="target_b"),
+            ],
+        )
+        # Manually set the selection function
+        group._selection_func = select_first_target
+
+        targets = route_message_through_edge_groups([group], "src", "test")
+
+        assert targets == ["target_a"]
+
+    def test_switch_case_without_selection_func_broadcasts(self) -> None:
+        """Test SwitchCaseEdgeGroup without selection_func broadcasts to all."""
+        group = SwitchCaseEdgeGroup(
+            source_id="src",
+            cases=[
+                SwitchCaseEdgeGroupCase(condition=lambda m: True, target_id="target_a"),
+                SwitchCaseEdgeGroupDefault(target_id="target_b"),
+            ],
+        )
+        group._selection_func = None
+
+        targets = route_message_through_edge_groups([group], "src", "test")
+
+        assert set(targets) == {"target_a", "target_b"}
+
+    def test_fan_out_with_selection_func(self) -> None:
+        """Test FanOutEdgeGroup uses selection_func."""
+
+        def select_all(msg: Any, targets: list[str]) -> list[str]:
+            return targets
+
+        group = FanOutEdgeGroup(
+            source_id="src",
+            target_ids=["fan_a", "fan_b", "fan_c"],
+            selection_func=select_all,
+        )
+
+        targets = route_message_through_edge_groups([group], "src", "broadcast")
+
+        assert set(targets) == {"fan_a", "fan_b", "fan_c"}
+
+    def test_fan_in_is_not_routed_directly(self) -> None:
+        """Test FanInEdgeGroup is handled separately (not routed here)."""
+        group = FanInEdgeGroup(
+            source_ids=["src_a", "src_b"],
+            target_id="aggregator",
+        )
+
+        # Fan-in should not add targets through this function
+        targets = route_message_through_edge_groups([group], "src_a", "message")
+
+        assert targets == []
+
+    def test_multiple_edge_groups_aggregated(self) -> None:
+        """Test that targets from multiple edge groups are aggregated."""
+        group1 = SingleEdgeGroup(source_id="src", target_id="t1", condition=lambda m: True)
+        group2 = SingleEdgeGroup(source_id="src", target_id="t2", condition=lambda m: True)
+
+        targets = route_message_through_edge_groups([group1, group2], "src", "msg")
+
+        assert set(targets) == {"t1", "t2"}
+
+
+class TestBuildAgentExecutorResponse:
+    """Test suite for build_agent_executor_response function."""
+
+    def test_builds_response_with_text(self) -> None:
+        """Test building response with plain text."""
+        response = build_agent_executor_response(
+            executor_id="my_executor",
+            response_text="Hello, world!",
+            structured_response=None,
+            previous_message="User input",
+        )
+
+        assert response.executor_id == "my_executor"
+        assert response.agent_response.text == "Hello, world!"
+        assert len(response.full_conversation) == 2  # User + Assistant
+
+    def test_builds_response_with_structured_response(self) -> None:
+        """Test building response with structured JSON response."""
+        structured = {"answer": 42, "reason": "because"}
+
+        response = build_agent_executor_response(
+            executor_id="calc",
+            response_text="Original text",
+            structured_response=structured,
+            previous_message="Calculate",
+        )
+
+        # Structured response overrides text
+        assert response.agent_response.text == json.dumps(structured)
+
+    def test_conversation_includes_previous_string_message(self) -> None:
+        """Test that string previous_message is included in conversation."""
+        response = build_agent_executor_response(
+            executor_id="exec",
+            response_text="Response",
+            structured_response=None,
+            previous_message="User said this",
+        )
+
+        assert len(response.full_conversation) == 2
+        assert response.full_conversation[0].role == "user"
+        assert response.full_conversation[0].text == "User said this"
+        assert response.full_conversation[1].role == "assistant"
+
+    def test_conversation_extends_previous_agent_executor_response(self) -> None:
+        """Test that previous AgentExecutorResponse's conversation is extended."""
+        # Create a previous response with conversation history
+        previous = AgentExecutorResponse(
+            executor_id="prev",
+            agent_response=AgentResponse(messages=[Message(role="assistant", contents=["Previous"])]),
+            full_conversation=[
+                Message(role="user", contents=["First"]),
+                Message(role="assistant", contents=["Previous"]),
+            ],
+        )
+
+        response = build_agent_executor_response(
+            executor_id="current",
+            response_text="Current response",
+            structured_response=None,
+            previous_message=previous,
+        )
+
+        # Should have 3 messages: First + Previous + Current
+        assert len(response.full_conversation) == 3
+        assert response.full_conversation[0].text == "First"
+        assert response.full_conversation[1].text == "Previous"
+        assert response.full_conversation[2].text == "Current response"
+
+
+class TestExtractMessageContent:
+    """Test suite for _extract_message_content function."""
+
+    def test_extract_from_string(self) -> None:
+        """Test extracting content from plain string."""
+        result = _extract_message_content("Hello, world!")
+
+        assert result == "Hello, world!"
+
+    def test_extract_from_agent_executor_response_with_text(self) -> None:
+        """Test extracting from AgentExecutorResponse with text."""
+        response = AgentExecutorResponse(
+            executor_id="exec",
+            agent_response=AgentResponse(messages=[Message(role="assistant", contents=["Response text"])]),
+            full_conversation=[Message(role="assistant", contents=["Response text"])],
+        )
+
+        result = _extract_message_content(response)
+
+        assert result == "Response text"
+
+    def test_extract_from_agent_executor_response_with_messages(self) -> None:
+        """Test extracting from AgentExecutorResponse with messages."""
+        response = AgentExecutorResponse(
+            executor_id="exec",
+            agent_response=AgentResponse(
+                messages=[
+                    Message(role="user", contents=["First"]),
+                    Message(role="assistant", contents=["Last message"]),
+                ]
+            ),
+            full_conversation=[
+                Message(role="user", contents=["First"]),
+                Message(role="assistant", contents=["Last message"]),
+            ],
+        )
+
+        result = _extract_message_content(response)
+
+        # AgentResponse.text concatenates all message texts
+        assert result == "FirstLast message"
+
+    def test_extract_from_agent_executor_request(self) -> None:
+        """Test extracting from AgentExecutorRequest."""
+        request = AgentExecutorRequest(
+            messages=[
+                Message(role="user", contents=["First"]),
+                Message(role="user", contents=["Last request"]),
+            ]
+        )
+
+        result = _extract_message_content(request)
+
+        assert result == "Last request"
+
+    def test_extract_from_dict_returns_empty(self) -> None:
+        """Test that dict messages return empty string (unexpected input)."""
+        msg_dict = {"messages": [{"text": "Hello"}]}
+
+        result = _extract_message_content(msg_dict)
+
+        assert result == ""
+
+    def test_extract_returns_empty_for_unknown_type(self) -> None:
+        """Test that unknown types return empty string."""
+        result = _extract_message_content(12345)
+
+        assert result == ""
+
+
+class TestEdgeGroupIntegration:
+    """Integration tests for edge group routing with realistic scenarios."""
+
+    def test_conditional_routing_by_message_type(self) -> None:
+        """Test routing based on message content/type."""
+
+        @dataclass
+        class SpamResult:
+            is_spam: bool
+            reason: str
+
+        def is_spam_condition(msg: Any) -> bool:
+            if isinstance(msg, SpamResult):
+                return msg.is_spam
+            return False
+
+        def is_not_spam_condition(msg: Any) -> bool:
+            if isinstance(msg, SpamResult):
+                return not msg.is_spam
+            return False
+
+        spam_group = SingleEdgeGroup(
+            source_id="detector",
+            target_id="spam_handler",
+            condition=is_spam_condition,
+        )
+        legit_group = SingleEdgeGroup(
+            source_id="detector",
+            target_id="email_handler",
+            condition=is_not_spam_condition,
+        )
+
+        # Test spam message
+        spam_msg = SpamResult(is_spam=True, reason="Suspicious content")
+        targets = route_message_through_edge_groups([spam_group, legit_group], "detector", spam_msg)
+        assert targets == ["spam_handler"]
+
+        # Test legitimate message
+        legit_msg = SpamResult(is_spam=False, reason="Clean")
+        targets = route_message_through_edge_groups([spam_group, legit_group], "detector", legit_msg)
+        assert targets == ["email_handler"]
+
+    def test_fan_out_to_multiple_workers(self) -> None:
+        """Test fan-out to multiple parallel workers."""
+
+        def select_all_workers(msg: Any, targets: list[str]) -> list[str]:
+            return targets
+
+        group = FanOutEdgeGroup(
+            source_id="coordinator",
+            target_ids=["worker_1", "worker_2", "worker_3"],
+            selection_func=select_all_workers,
+        )
+
+        targets = route_message_through_edge_groups([group], "coordinator", {"task": "process"})
+
+        assert len(targets) == 3
+        assert set(targets) == {"worker_1", "worker_2", "worker_3"}