agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: Move Workflow, Edge, EdgeGroup and Executor to AFBaseModel (#472)

Browse Source 
* Refactor workflow to introduce EdgeRunner for edge execution.

* Fix edge cases

* Convert Workflow, Edge, EdgeGroup, and Executor into AFBaseModel to support object model serialization

* format

* remove accidental file

* fix typing

* Add type information to EdgeGroup and Executor subclasses

* fix format

* Add condition_name field to Edge

* Add new fields

* remove Optional

* Update
This commit is contained in:
Eric Zhu
2025-08-22 15:02:34 -07:00
committed by GitHub
Unverified
parent 88c51013c6
commit 6e9d35830f
13 changed files with 1588 additions and 633 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/workflow/tests/test_edge.py
+209 -201
View File
@@ -8,14 +8,17 @@ import pytest
from agent_framework.workflow import Executor, WorkflowContext, handler
from agent_framework_workflow._edge import (
Case,
Default,
Edge,
FanInEdgeGroup,
FanOutEdgeGroup,
SingleEdgeGroup,
SwitchCaseEdgeGroup,
SwitchCaseEdgeGroupCase,
SwitchCaseEdgeGroupDefault,
)
from agent_framework_workflow._edge_runner import create_edge_runner
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
@dataclass
@@ -35,28 +38,40 @@ class MockMessageSecondary:
class MockExecutor(Executor):
"""A mock executor for testing purposes."""
call_count: int = 0
last_message: Any = None
@handler
async def mock_handler(self, message: MockMessage, ctx: WorkflowContext) -> None:
"""A mock handler that does nothing."""
pass
self.call_count += 1
self.last_message = message
class MockExecutorSecondary(Executor):
"""A secondary mock executor for testing purposes."""
call_count: int = 0
last_message: Any = None
@handler
async def mock_handler_secondary(self, message: MockMessageSecondary, ctx: WorkflowContext) -> None:
"""A secondary mock handler that does nothing."""
pass
self.call_count += 1
self.last_message = message
class MockAggregator(Executor):
"""A mock aggregator for testing purposes."""
call_count: int = 0
last_message: Any = None
@handler
async def mock_aggregator_handler(self, message: list[MockMessage], ctx: WorkflowContext) -> None:
"""A mock aggregator handler that does nothing."""
pass
self.call_count += 1
self.last_message = message
# region Edge
@@ -67,7 +82,7 @@ def test_create_edge():
source = MockExecutor(id="source_executor")
target = MockExecutor(id="target_executor")
edge = Edge(source=source, target=target)
edge = Edge(source_id=source.id, target_id=target.id)
assert edge.source_id == "source_executor"
assert edge.target_id == "target_executor"
@@ -79,9 +94,9 @@ def test_edge_can_handle():
source = MockExecutor(id="source_executor")
target = MockExecutor(id="target_executor")
edge = Edge(source=source, target=target)
edge = Edge(source_id=source.id, target_id=target.id)
assert edge.can_handle(MockMessage(data="test"))
assert edge.should_route(MockMessage(data="test"))
# endregion Edge
@@ -94,10 +109,10 @@ def test_single_edge_group():
source = MockExecutor(id="source_executor")
target = MockExecutor(id="target_executor")
edge_group = SingleEdgeGroup(source=source, target=target)
edge_group = SingleEdgeGroup(source_id=source.id, target_id=target.id)
assert edge_group.source_executors == [source]
assert edge_group.target_executors == [target]
assert edge_group.source_executor_ids == [source.id]
assert edge_group.target_executor_ids == [target.id]
assert edge_group.edges[0].source_id == "source_executor"
assert edge_group.edges[0].target_id == "target_executor"
@@ -107,92 +122,88 @@ def test_single_edge_group_with_condition():
source = MockExecutor(id="source_executor")
target = MockExecutor(id="target_executor")
edge_group = SingleEdgeGroup(source=source, target=target, condition=lambda x: x.data == "test")
edge_group = SingleEdgeGroup(source_id=source.id, target_id=target.id, condition=lambda x: x.data == "test")
assert edge_group.source_executors == [source]
assert edge_group.target_executors == [target]
assert edge_group.source_executor_ids == [source.id]
assert edge_group.target_executor_ids == [target.id]
assert edge_group.edges[0].source_id == "source_executor"
assert edge_group.edges[0].target_id == "target_executor"
assert edge_group.edges[0]._condition is not None # type: ignore
async def test_single_edge_group_send_message():
"""Test sending a message through a single edge group."""
async def test_single_edge_group_send_message() -> None:
"""Test sending a message through a single edge runner."""
source = MockExecutor(id="source_executor")
target = MockExecutor(id="target_executor")
edge_group = SingleEdgeGroup(source=source, target=target)
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
executors: dict[str, Executor] = {source.id: source, target.id: target}
edge_group = SingleEdgeGroup(source_id=source.id, target_id=target.id)
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
data = MockMessage(data="test")
message = Message(data=data, source_id=source.id)
success = await edge_group.send_message(message, shared_state, ctx)
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is True
async def test_single_edge_group_send_message_with_target():
"""Test sending a message through a single edge group."""
async def test_single_edge_group_send_message_with_target() -> None:
"""Test sending a message through a single edge runner."""
source = MockExecutor(id="source_executor")
target = MockExecutor(id="target_executor")
edge_group = SingleEdgeGroup(source=source, target=target)
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
executors: dict[str, Executor] = {source.id: source, target.id: target}
edge_group = SingleEdgeGroup(source_id=source.id, target_id=target.id)
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
data = MockMessage(data="test")
message = Message(data=data, source_id=source.id, target_id=target.id)
success = await edge_group.send_message(message, shared_state, ctx)
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is True
async def test_single_edge_group_send_message_with_invalid_target():
"""Test sending a message through a single edge group."""
async def test_single_edge_group_send_message_with_invalid_target() -> None:
"""Test sending a message through a single edge runner."""
source = MockExecutor(id="source_executor")
target = MockExecutor(id="target_executor")
edge_group = SingleEdgeGroup(source=source, target=target)
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
executors: dict[str, Executor] = {source.id: source, target.id: target}
edge_group = SingleEdgeGroup(source_id=source.id, target_id=target.id)
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
data = MockMessage(data="test")
message = Message(data=data, source_id=source.id, target_id="invalid_target")
success = await edge_group.send_message(message, shared_state, ctx)
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is False
async def test_single_edge_group_send_message_with_invalid_data():
"""Test sending a message through a single edge group."""
async def test_single_edge_group_send_message_with_invalid_data() -> None:
"""Test sending a message through a single edge runner with invalid data."""
source = MockExecutor(id="source_executor")
target = MockExecutor(id="target_executor")
edge_group = SingleEdgeGroup(source=source, target=target)
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
executors: dict[str, Executor] = {source.id: source, target.id: target}
edge_group = SingleEdgeGroup(source_id=source.id, target_id=target.id)
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
data = "invalid_data"
message = Message(data=data, source_id=source.id)
success = await edge_group.send_message(message, shared_state, ctx)
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is False
@@ -208,10 +219,10 @@ def test_source_edge_group():
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutor(id="target_executor_2")
edge_group = FanOutEdgeGroup(source=source, targets=[target1, target2])
edge_group = FanOutEdgeGroup(source_id=source.id, target_ids=[target1.id, target2.id])
assert edge_group.source_executors == [source]
assert edge_group.target_executors == [target1, target2]
assert edge_group.source_executor_ids == [source.id]
assert edge_group.target_executor_ids == [target1.id, target2.id]
assert len(edge_group.edges) == 2
assert edge_group.edges[0].source_id == "source_executor"
assert edge_group.edges[0].target_id == "target_executor_1"
@@ -219,128 +230,122 @@ def test_source_edge_group():
assert edge_group.edges[1].target_id == "target_executor_2"
def test_source_edge_group_invalid_number_of_targets():
def test_source_edge_group_invalid_number_of_targets() -> None:
"""Test creating a fan-out group with an invalid number of targets."""
source = MockExecutor(id="source_executor")
target = MockExecutor(id="target_executor")
with pytest.raises(ValueError, match="FanOutEdgeGroup must contain at least two targets"):
FanOutEdgeGroup(source=source, targets=[target])
FanOutEdgeGroup(source_id=source.id, target_ids=[target.id])
async def test_source_edge_group_send_message():
"""Test sending a message through a fan-out group."""
async def test_source_edge_group_send_message() -> None:
"""Test sending a message through a fan-out edge runner."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutor(id="target_executor_2")
edge_group = FanOutEdgeGroup(source=source, targets=[target1, target2])
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2}
edge_group = FanOutEdgeGroup(source_id=source.id, target_ids=[target1.id, target2.id])
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
data = MockMessage(data="test")
message = Message(data=data, source_id=source.id)
with patch("agent_framework_workflow._edge.Edge.send_message") as mock_send:
success = await edge_group.send_message(message, shared_state, ctx)
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is True
assert mock_send.call_count == 2
assert success is True
assert target1.call_count == 1
assert target2.call_count == 1
async def test_source_edge_group_send_message_with_target():
async def test_source_edge_group_send_message_with_target() -> None:
"""Test sending a message through a fan-out group with a target."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutor(id="target_executor_2")
edge_group = FanOutEdgeGroup(source=source, targets=[target1, target2])
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
edge_group = FanOutEdgeGroup(source_id=source.id, target_ids=[target1.id, target2.id])
executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2}
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
data = MockMessage(data="test")
message = Message(data=data, source_id=source.id, target_id=target1.id)
with patch("agent_framework_workflow._edge.Edge.send_message") as mock_send:
success = await edge_group.send_message(message, shared_state, ctx)
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is True
assert mock_send.call_count == 1
assert mock_send.call_args[0][0].target_id == target1.id
assert success is True
assert target1.call_count == 1
assert target2.call_count == 0 # target2 should not be called since message targets target1
async def test_source_edge_group_send_message_with_invalid_target():
async def test_source_edge_group_send_message_with_invalid_target() -> None:
"""Test sending a message through a fan-out group with an invalid target."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutor(id="target_executor_2")
edge_group = FanOutEdgeGroup(source=source, targets=[target1, target2])
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
edge_group = FanOutEdgeGroup(source_id=source.id, target_ids=[target1.id, target2.id])
executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2}
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
data = MockMessage(data="test")
message = Message(data=data, source_id=source.id, target_id="invalid_target")
success = await edge_group.send_message(message, shared_state, ctx)
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is False
async def test_source_edge_group_send_message_with_invalid_data():
async def test_source_edge_group_send_message_with_invalid_data() -> None:
"""Test sending a message through a fan-out group with invalid data."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutor(id="target_executor_2")
edge_group = FanOutEdgeGroup(source=source, targets=[target1, target2])
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
edge_group = FanOutEdgeGroup(source_id=source.id, target_ids=[target1.id, target2.id])
executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2}
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
data = "invalid_data"
message = Message(data=data, source_id=source.id)
success = await edge_group.send_message(message, shared_state, ctx)
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is False
async def test_source_edge_group_send_message_only_one_successful_send():
async def test_source_edge_group_send_message_only_one_successful_send() -> None:
"""Test sending a message through a fan-out group where only one edge can handle the message."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutorSecondary(id="target_executor_2")
edge_group = FanOutEdgeGroup(source=source, targets=[target1, target2])
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
edge_group = FanOutEdgeGroup(source_id=source.id, target_ids=[target1.id, target2.id])
executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2}
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
data = MockMessage(data="test")
message = Message(data=data, source_id=source.id)
with patch("agent_framework_workflow._edge.Edge.send_message") as mock_send:
success = await edge_group.send_message(message, shared_state, ctx)
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is True
assert mock_send.call_count == 1
assert success is True
assert target1.call_count == 1 # target1 can handle MockMessage
assert target2.call_count == 0 # target2 (MockExecutorSecondary) cannot handle MockMessage
def test_source_edge_group_with_selection_func():
@@ -350,13 +355,13 @@ def test_source_edge_group_with_selection_func():
target2 = MockExecutor(id="target_executor_2")
edge_group = FanOutEdgeGroup(
source=source,
targets=[target1, target2],
source_id=source.id,
target_ids=[target1.id, target2.id],
selection_func=lambda data, target_ids: [target1.id],
)
assert edge_group.source_executors == [source]
assert edge_group.target_executors == [target1, target2]
assert edge_group.source_executor_ids == [source.id]
assert edge_group.target_executor_ids == [target1.id, target2.id]
assert len(edge_group.edges) == 2
assert edge_group.edges[0].source_id == "source_executor"
assert edge_group.edges[0].target_id == "target_executor_1"
@@ -364,20 +369,20 @@ def test_source_edge_group_with_selection_func():
assert edge_group.edges[1].target_id == "target_executor_2"
async def test_source_edge_group_with_selection_func_send_message():
async def test_source_edge_group_with_selection_func_send_message() -> None:
"""Test sending a message through a fan-out group with a selection function."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutor(id="target_executor_2")
edge_group = FanOutEdgeGroup(
source=source,
targets=[target1, target2],
source_id=source.id,
target_ids=[target1.id, target2.id],
selection_func=lambda data, target_ids: [target1.id, target2.id],
)
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2}
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
@@ -385,28 +390,28 @@ async def test_source_edge_group_with_selection_func_send_message():
data = MockMessage(data="test")
message = Message(data=data, source_id=source.id)
with patch("agent_framework_workflow._edge.Edge.send_message") as mock_send:
success = await edge_group.send_message(message, shared_state, ctx)
with patch("agent_framework_workflow._edge_runner.EdgeRunner._execute_on_target") as mock_send:
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is True
assert mock_send.call_count == 2
async def test_source_edge_group_with_selection_func_send_message_with_invalid_selection_result():
async def test_source_edge_group_with_selection_func_send_message_with_invalid_selection_result() -> None:
"""Test sending a message through a fan-out group with a selection func with an invalid selection result."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutor(id="target_executor_2")
edge_group = FanOutEdgeGroup(
source=source,
targets=[target1, target2],
source_id=source.id,
target_ids=[target1.id, target2.id],
selection_func=lambda data, target_ids: [target1.id, "invalid_target"],
)
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2}
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
@@ -414,23 +419,23 @@ async def test_source_edge_group_with_selection_func_send_message_with_invalid_s
message = Message(data=data, source_id=source.id)
with pytest.raises(RuntimeError):
await edge_group.send_message(message, shared_state, ctx)
await edge_runner.send_message(message, shared_state, ctx)
async def test_source_edge_group_with_selection_func_send_message_with_target():
async def test_source_edge_group_with_selection_func_send_message_with_target() -> None:
"""Test sending a message through a fan-out group with a selection func with a target."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutor(id="target_executor_2")
edge_group = FanOutEdgeGroup(
source=source,
targets=[target1, target2],
source_id=source.id,
target_ids=[target1.id, target2.id],
selection_func=lambda data, target_ids: [target1.id, target2.id],
)
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2}
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
@@ -438,28 +443,28 @@ async def test_source_edge_group_with_selection_func_send_message_with_target():
data = MockMessage(data="test")
message = Message(data=data, source_id=source.id, target_id=target1.id)
with patch("agent_framework_workflow._edge.Edge.send_message") as mock_send:
success = await edge_group.send_message(message, shared_state, ctx)
with patch("agent_framework_workflow._edge_runner.EdgeRunner._execute_on_target") as mock_send:
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is True
assert mock_send.call_count == 1
assert mock_send.call_args[0][0].target_id == target1.id
assert mock_send.call_args[0][0] == target1.id
async def test_source_edge_group_with_selection_func_send_message_with_target_not_in_selection():
async def test_source_edge_group_with_selection_func_send_message_with_target_not_in_selection() -> None:
"""Test sending a message through a fan-out group with a selection func with a target not in the selection."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutor(id="target_executor_2")
edge_group = FanOutEdgeGroup(
source=source,
targets=[target1, target2],
source_id=source.id,
target_ids=[target1.id, target2.id],
selection_func=lambda data, target_ids: [target1.id], # Only target1 will receive the message
)
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2}
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
@@ -467,22 +472,24 @@ async def test_source_edge_group_with_selection_func_send_message_with_target_no
data = MockMessage(data="test")
message = Message(data=data, source_id=source.id, target_id=target2.id)
success = await edge_group.send_message(message, shared_state, ctx)
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is False
async def test_source_edge_group_with_selection_func_send_message_with_invalid_data():
async def test_source_edge_group_with_selection_func_send_message_with_invalid_data() -> None:
"""Test sending a message through a fan-out group with a selection func with invalid data."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutor(id="target_executor_2")
edge_group = FanOutEdgeGroup(
source=source, targets=[target1, target2], selection_func=lambda data, target_ids: [target1.id, target2.id]
source_id=source.id,
target_ids=[target1.id, target2.id],
selection_func=lambda data, target_ids: [target1.id, target2.id],
)
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2}
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
@@ -490,22 +497,24 @@ async def test_source_edge_group_with_selection_func_send_message_with_invalid_d
data = "invalid_data"
message = Message(data=data, source_id=source.id)
success = await edge_group.send_message(message, shared_state, ctx)
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is False
async def test_source_edge_group_with_selection_func_send_message_with_target_invalid_data():
async def test_source_edge_group_with_selection_func_send_message_with_target_invalid_data() -> None:
"""Test sending a message through a fan-out group with a selection func with a target and invalid data."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutor(id="target_executor_2")
edge_group = FanOutEdgeGroup(
source=source, targets=[target1, target2], selection_func=lambda data, target_ids: [target1.id, target2.id]
source_id=source.id,
target_ids=[target1.id, target2.id],
selection_func=lambda data, target_ids: [target1.id, target2.id],
)
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2}
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
@@ -513,7 +522,7 @@ async def test_source_edge_group_with_selection_func_send_message_with_target_in
data = "invalid_data"
message = Message(data=data, source_id=source.id, target_id=target1.id)
success = await edge_group.send_message(message, shared_state, ctx)
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is False
@@ -528,10 +537,10 @@ def test_target_edge_group():
source2 = MockExecutor(id="source_executor_2")
target = MockAggregator(id="target_executor")
edge_group = FanInEdgeGroup(sources=[source1, source2], target=target)
edge_group = FanInEdgeGroup(source_ids=[source1.id, source2.id], target_id=target.id)
assert edge_group.source_executors == [source1, source2]
assert edge_group.target_executors == [target]
assert edge_group.source_executor_ids == [source1.id, source2.id]
assert edge_group.target_executor_ids == [target.id]
assert len(edge_group.edges) == 2
assert edge_group.edges[0].source_id == "source_executor_1"
assert edge_group.edges[0].target_id == "target_executor"
@@ -545,27 +554,27 @@ def test_target_edge_group_invalid_number_of_sources():
target = MockAggregator(id="target_executor")
with pytest.raises(ValueError, match="FanInEdgeGroup must contain at least two sources"):
FanInEdgeGroup(sources=[source], target=target)
FanInEdgeGroup(source_ids=[source.id], target_id=target.id)
async def test_target_edge_group_send_message_buffer():
async def test_target_edge_group_send_message_buffer() -> None:
"""Test sending a message through a fan-in edge group with buffering."""
source1 = MockExecutor(id="source_executor_1")
source2 = MockExecutor(id="source_executor_2")
target = MockAggregator(id="target_executor")
edge_group = FanInEdgeGroup(sources=[source1, source2], target=target)
edge_group = FanInEdgeGroup(source_ids=[source1.id, source2.id], target_id=target.id)
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
executors: dict[str, Executor] = {source1.id: source1, source2.id: source2, target.id: target}
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
data = MockMessage(data="test")
with patch("agent_framework_workflow._edge.Edge.send_message") as mock_send:
success = await edge_group.send_message(
with patch("agent_framework_workflow._edge_runner.EdgeRunner._execute_on_target") as mock_send:
success = await edge_runner.send_message(
Message(data=data, source_id=source1.id),
shared_state,
ctx,
@@ -573,9 +582,9 @@ async def test_target_edge_group_send_message_buffer():
assert success is True
assert mock_send.call_count == 0 # The message should be buffered and wait for the second source
assert len(edge_group._buffer[source1.id]) == 1 # type: ignore
assert len(edge_runner._buffer[source1.id]) == 1 # type: ignore
success = await edge_group.send_message(
success = await edge_runner.send_message(
Message(data=data, source_id=source2.id),
shared_state,
ctx,
@@ -584,19 +593,19 @@ async def test_target_edge_group_send_message_buffer():
assert mock_send.call_count == 1 # The message should be sent now that both sources have sent their messages
# Buffer should be cleared after sending
assert not edge_group._buffer # type: ignore
assert not edge_runner._buffer # type: ignore
async def test_target_edge_group_send_message_with_invalid_target():
async def test_target_edge_group_send_message_with_invalid_target() -> None:
"""Test sending a message through a fan-in edge group with an invalid target."""
source1 = MockExecutor(id="source_executor_1")
source2 = MockExecutor(id="source_executor_2")
target = MockAggregator(id="target_executor")
edge_group = FanInEdgeGroup(sources=[source1, source2], target=target)
edge_group = FanInEdgeGroup(source_ids=[source1.id, source2.id], target_id=target.id)
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
executors: dict[str, Executor] = {source1.id: source1, source2.id: source2, target.id: target}
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
@@ -604,20 +613,20 @@ async def test_target_edge_group_send_message_with_invalid_target():
data = MockMessage(data="test")
message = Message(data=data, source_id=source1.id, target_id="invalid_target")
success = await edge_group.send_message(message, shared_state, ctx)
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is False
async def test_target_edge_group_send_message_with_invalid_data():
async def test_target_edge_group_send_message_with_invalid_data() -> None:
"""Test sending a message through a fan-in edge group with invalid data."""
source1 = MockExecutor(id="source_executor_1")
source2 = MockExecutor(id="source_executor_2")
target = MockAggregator(id="target_executor")
edge_group = FanInEdgeGroup(sources=[source1, source2], target=target)
edge_group = FanInEdgeGroup(source_ids=[source1.id, source2.id], target_id=target.id)
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
executors: dict[str, Executor] = {source1.id: source1, source2.id: source2, target.id: target}
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
@@ -625,7 +634,7 @@ async def test_target_edge_group_send_message_with_invalid_data():
data = "invalid_data"
message = Message(data=data, source_id=source1.id)
success = await edge_group.send_message(message, shared_state, ctx)
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is False
@@ -634,22 +643,22 @@ async def test_target_edge_group_send_message_with_invalid_data():
# region SwitchCaseEdgeGroup
def test_switch_case_edge_group():
def test_switch_case_edge_group() -> None:
"""Test creating a switch case edge group."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutor(id="target_executor_2")
edge_group = SwitchCaseEdgeGroup(
source=source,
source_id=source.id,
cases=[
Case(condition=lambda x: x.data < 0, target=target1),
Default(target=target2),
SwitchCaseEdgeGroupCase(condition=lambda x: x.data < 0, target_id=target1.id),
SwitchCaseEdgeGroupDefault(target_id=target2.id),
],
)
assert edge_group.source_executors == [source]
assert edge_group.target_executors == [target1, target2]
assert edge_group.source_executor_ids == [source.id]
assert edge_group.target_executor_ids == [target1.id, target2.id]
assert len(edge_group.edges) == 2
assert edge_group.edges[0].source_id == "source_executor"
assert edge_group.edges[0].target_id == "target_executor_1"
@@ -670,18 +679,18 @@ def test_switch_case_edge_group_invalid_number_of_cases():
ValueError, match=r"SwitchCaseEdgeGroup must contain at least two cases \(including the default case\)."
):
SwitchCaseEdgeGroup(
source=source,
source_id=source.id,
cases=[
Case(condition=lambda x: x.data < 0, target=target),
SwitchCaseEdgeGroupCase(condition=lambda x: x.data < 0, target_id=target.id),
],
)
with pytest.raises(ValueError, match="SwitchCaseEdgeGroup must contain exactly one default case."):
SwitchCaseEdgeGroup(
source=source,
source_id=source.id,
cases=[
Case(condition=lambda x: x.data < 0, target=target),
Case(condition=lambda x: x.data >= 0, target=target),
SwitchCaseEdgeGroupCase(condition=lambda x: x.data < 0, target_id=target.id),
SwitchCaseEdgeGroupCase(condition=lambda x: x.data >= 0, target_id=target.id),
],
)
@@ -694,31 +703,30 @@ def test_switch_case_edge_group_invalid_number_of_default_cases():
with pytest.raises(ValueError, match="SwitchCaseEdgeGroup must contain exactly one default case."):
SwitchCaseEdgeGroup(
source=source,
source_id=source.id,
cases=[
Case(condition=lambda x: x.data < 0, target=target1),
Default(target=target2),
Default(target=target2),
SwitchCaseEdgeGroupCase(condition=lambda x: x.data < 0, target_id=target1.id),
SwitchCaseEdgeGroupDefault(target_id=target2.id),
SwitchCaseEdgeGroupDefault(target_id=target2.id),
],
)
async def test_switch_case_edge_group_send_message():
async def test_switch_case_edge_group_send_message() -> None:
"""Test sending a message through a switch case edge group."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutor(id="target_executor_2")
edge_group = SwitchCaseEdgeGroup(
source=source,
source_id=source.id,
cases=[
Case(condition=lambda x: x.data < 0, target=target1),
Default(target=target2),
SwitchCaseEdgeGroupCase(condition=lambda x: x.data < 0, target_id=target1.id),
SwitchCaseEdgeGroupDefault(target_id=target2.id),
],
)
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2}
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
@@ -726,8 +734,8 @@ async def test_switch_case_edge_group_send_message():
data = MockMessage(data=-1)
message = Message(data=data, source_id=source.id)
with patch("agent_framework_workflow._edge.Edge.send_message") as mock_send:
success = await edge_group.send_message(message, shared_state, ctx)
with patch("agent_framework_workflow._edge_runner.EdgeRunner._execute_on_target") as mock_send:
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is True
assert mock_send.call_count == 1
@@ -735,29 +743,29 @@ async def test_switch_case_edge_group_send_message():
# Default condition should
data = MockMessage(data=1)
message = Message(data=data, source_id=source.id)
with patch("agent_framework_workflow._edge.Edge.send_message") as mock_send:
success = await edge_group.send_message(message, shared_state, ctx)
with patch("agent_framework_workflow._edge_runner.EdgeRunner._execute_on_target") as mock_send:
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is True
assert mock_send.call_count == 1
async def test_switch_case_edge_group_send_message_with_invalid_target():
async def test_switch_case_edge_group_send_message_with_invalid_target() -> None:
"""Test sending a message through a switch case edge group with an invalid target."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutor(id="target_executor_2")
edge_group = SwitchCaseEdgeGroup(
source=source,
source_id=source.id,
cases=[
Case(condition=lambda x: x.data < 0, target=target1),
Default(target=target2),
SwitchCaseEdgeGroupCase(condition=lambda x: x.data < 0, target_id=target1.id),
SwitchCaseEdgeGroupDefault(target_id=target2.id),
],
)
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2}
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
@@ -765,26 +773,26 @@ async def test_switch_case_edge_group_send_message_with_invalid_target():
data = MockMessage(data=-1)
message = Message(data=data, source_id=source.id, target_id="invalid_target")
success = await edge_group.send_message(message, shared_state, ctx)
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is False
async def test_switch_case_edge_group_send_message_with_valid_target():
async def test_switch_case_edge_group_send_message_with_valid_target() -> None:
"""Test sending a message through a switch case edge group with a target."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutor(id="target_executor_2")
edge_group = SwitchCaseEdgeGroup(
source=source,
source_id=source.id,
cases=[
Case(condition=lambda x: x.data < 0, target=target1),
Default(target=target2),
SwitchCaseEdgeGroupCase(condition=lambda x: x.data < 0, target_id=target1.id),
SwitchCaseEdgeGroupDefault(target_id=target2.id),
],
)
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2}
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
@@ -792,31 +800,31 @@ async def test_switch_case_edge_group_send_message_with_valid_target():
data = MockMessage(data=1) # Condition will fail
message = Message(data=data, source_id=source.id, target_id=target1.id)
success = await edge_group.send_message(message, shared_state, ctx)
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is False
data = MockMessage(data=-1) # Condition will pass
message = Message(data=data, source_id=source.id, target_id=target1.id)
success = await edge_group.send_message(message, shared_state, ctx)
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is True
async def test_switch_case_edge_group_send_message_with_invalid_data():
async def test_switch_case_edge_group_send_message_with_invalid_data() -> None:
"""Test sending a message through a switch case edge group with invalid data."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutor(id="target_executor_2")
edge_group = SwitchCaseEdgeGroup(
source=source,
source_id=source.id,
cases=[
Case(condition=lambda x: x.data < 0, target=target1),
Default(target=target2),
SwitchCaseEdgeGroupCase(condition=lambda x: x.data < 0, target_id=target1.id),
SwitchCaseEdgeGroupDefault(target_id=target2.id),
],
)
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
from agent_framework_workflow._shared_state import SharedState
executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2}
edge_runner = create_edge_runner(edge_group, executors)
shared_state = SharedState()
ctx = InProcRunnerContext()
@@ -824,7 +832,7 @@ async def test_switch_case_edge_group_send_message_with_invalid_data():
data = "invalid_data"
message = Message(data=data, source_id=source.id)
success = await edge_group.send_message(message, shared_state, ctx)
success = await edge_runner.send_message(message, shared_state, ctx)
assert success is False
python/packages/workflow/tests/test_runner.py
+17 -12
View File
@@ -37,11 +37,13 @@ def test_create_runner():
# Create a loop
edge_groups = [
SingleEdgeGroup(executor_a, executor_b),
SingleEdgeGroup(executor_b, executor_a),
SingleEdgeGroup(executor_a.id, executor_b.id),
SingleEdgeGroup(executor_b.id, executor_a.id),
]
runner = Runner(edge_groups, shared_state=SharedState(), ctx=InProcRunnerContext())
executors: dict[str, Executor] = {executor_a.id: executor_a, executor_b.id: executor_b}
runner = Runner(edge_groups, executors, shared_state=SharedState(), ctx=InProcRunnerContext())
assert runner.context is not None and isinstance(runner.context, RunnerContext)
@@ -53,14 +55,15 @@ async def test_runner_run_until_convergence():
# Create a loop
edges = [
SingleEdgeGroup(executor_a, executor_b),
SingleEdgeGroup(executor_b, executor_a),
SingleEdgeGroup(executor_a.id, executor_b.id),
SingleEdgeGroup(executor_b.id, executor_a.id),
]
executors: dict[str, Executor] = {executor_a.id: executor_a, executor_b.id: executor_b}
shared_state = SharedState()
ctx = InProcRunnerContext()
runner = Runner(edges, shared_state, ctx)
runner = Runner(edges, executors, shared_state, ctx)
result: int | None = None
await executor_a.execute(
@@ -87,14 +90,15 @@ async def test_runner_run_until_convergence_not_completed():
# Create a loop
edges = [
SingleEdgeGroup(executor_a, executor_b),
SingleEdgeGroup(executor_b, executor_a),
SingleEdgeGroup(executor_a.id, executor_b.id),
SingleEdgeGroup(executor_b.id, executor_a.id),
]
executors: dict[str, Executor] = {executor_a.id: executor_a, executor_b.id: executor_b}
shared_state = SharedState()
ctx = InProcRunnerContext()
runner = Runner(edges, shared_state, ctx, max_iterations=5)
runner = Runner(edges, executors, shared_state, ctx, max_iterations=5)
await executor_a.execute(
MockMessage(data=0),
@@ -117,14 +121,15 @@ async def test_runner_already_running():
# Create a loop
edges = [
SingleEdgeGroup(executor_a, executor_b),
SingleEdgeGroup(executor_b, executor_a),
SingleEdgeGroup(executor_a.id, executor_b.id),
SingleEdgeGroup(executor_b.id, executor_a.id),
]
executors: dict[str, Executor] = {executor_a.id: executor_a, executor_b.id: executor_b}
shared_state = SharedState()
ctx = InProcRunnerContext()
runner = Runner(edges, shared_state, ctx)
runner = Runner(edges, executors, shared_state, ctx)
await executor_a.execute(
MockMessage(data=0),
python/packages/workflow/tests/test_serialization.py
+636
View File
@@ -0,0 +1,636 @@
# Copyright (c) Microsoft. All rights reserved.
import json
from typing import Any
import pytest
from agent_framework.workflow import Executor, WorkflowBuilder, WorkflowContext, handler
from agent_framework_workflow._edge import (
Edge,
FanInEdgeGroup,
FanOutEdgeGroup,
SingleEdgeGroup,
SwitchCaseEdgeGroup,
SwitchCaseEdgeGroupCase,
SwitchCaseEdgeGroupDefault,
)
class SampleExecutor(Executor):
"""Sample executor for serialization testing."""
@handler
async def handle_str(self, message: str, ctx: WorkflowContext[str]) -> None:
"""Handle string messages."""
await ctx.send_message(f"Processed: {message}")
class SampleAggregator(Executor):
"""Sample aggregator executor that can handle lists of messages."""
@handler
async def handle_str_list(self, messages: list[str], ctx: WorkflowContext[str]) -> None:
"""Handle list of string messages for fan-in aggregation."""
combined = " | ".join(messages)
await ctx.send_message(f"Aggregated: {combined}")
class TestSerializationWorkflowClasses:
"""Test serialization of workflow classes."""
def test_executor_serialization(self) -> None:
"""Test that Executor can be serialized and has correct fields, including type."""
executor = SampleExecutor(id="test-executor")
# Test model_dump
data = executor.model_dump()
assert data["id"] == "test-executor"
# Test type field
assert "type" in data, "Executor should have 'type' field"
assert data["type"] == "SampleExecutor", f"Expected type 'SampleExecutor', got {data['type']}"
# Test model_dump_json
json_str = executor.model_dump_json()
parsed = json.loads(json_str)
assert parsed["id"] == "test-executor"
# Test type field in JSON
assert "type" in parsed, "JSON should have 'type' field"
assert parsed["type"] == "SampleExecutor", "JSON should preserve type field"
def test_edge_serialization(self) -> None:
"""Test that Edge can be serialized and has correct fields."""
# Test edge without condition
edge = Edge(source_id="source", target_id="target")
# Test model_dump
data = edge.model_dump()
assert data["source_id"] == "source"
assert data["target_id"] == "target"
assert "condition_name" not in data or data["condition_name"] is None
# Test model_dump_json
json_str = edge.model_dump_json()
parsed = json.loads(json_str)
assert parsed["source_id"] == "source"
assert parsed["target_id"] == "target"
assert "condition_name" not in parsed or parsed["condition_name"] is None
def test_edge_serialization_with_named_condition(self) -> None:
"""Test that Edge with named function condition serializes condition_name correctly."""
def is_positive(x: int) -> bool:
return x > 0
edge = Edge(source_id="source", target_id="target", condition=is_positive)
# Test model_dump
data = edge.model_dump()
assert data["source_id"] == "source"
assert data["target_id"] == "target"
assert data["condition_name"] == "is_positive"
# Test model_dump_json
json_str = edge.model_dump_json()
parsed = json.loads(json_str)
assert parsed["source_id"] == "source"
assert parsed["target_id"] == "target"
assert parsed["condition_name"] == "is_positive"
def test_edge_serialization_with_lambda_condition(self) -> None:
"""Test that Edge with lambda condition serializes condition_name as '<lambda>'."""
edge = Edge(source_id="source", target_id="target", condition=lambda x: x > 0)
# Test model_dump
data = edge.model_dump()
assert data["source_id"] == "source"
assert data["target_id"] == "target"
assert data["condition_name"] == "<lambda>"
# Test model_dump_json
json_str = edge.model_dump_json()
parsed = json.loads(json_str)
assert parsed["source_id"] == "source"
assert parsed["target_id"] == "target"
assert parsed["condition_name"] == "<lambda>"
def test_single_edge_group_serialization(self) -> None:
"""Test that SingleEdgeGroup can be serialized and has correct fields, including edges and type."""
edge_group = SingleEdgeGroup(source_id="source", target_id="target")
# Test model_dump
data = edge_group.model_dump()
assert "id" in data
assert data["id"].startswith("SingleEdgeGroup/")
# Test type field
assert "type" in data, "SingleEdgeGroup should have 'type' field"
assert data["type"] == "SingleEdgeGroup", f"Expected type 'SingleEdgeGroup', got {data['type']}"
# Verify edges field is present and contains the edge
assert "edges" in data, "SingleEdgeGroup should have 'edges' field"
assert len(data["edges"]) == 1, "SingleEdgeGroup should have exactly one edge"
edge = data["edges"][0]
assert "source_id" in edge, "Edge should have source_id"
assert "target_id" in edge, "Edge should have target_id"
assert edge["source_id"] == "source", f"Expected source_id 'source', got {edge['source_id']}"
assert edge["target_id"] == "target", f"Expected target_id 'target', got {edge['target_id']}"
# Test model_dump_json
json_str = edge_group.model_dump_json()
parsed = json.loads(json_str)
assert "id" in parsed
assert parsed["id"].startswith("SingleEdgeGroup/")
# Test type field in JSON
assert "type" in parsed, "JSON should have 'type' field"
assert parsed["type"] == "SingleEdgeGroup", "JSON should preserve type field"
# Verify edges are preserved in JSON
assert "edges" in parsed, "JSON should have 'edges' field"
assert len(parsed["edges"]) == 1, "JSON should have exactly one edge"
json_edge = parsed["edges"][0]
assert json_edge["source_id"] == "source", "JSON should preserve edge source_id"
assert json_edge["target_id"] == "target", "JSON should preserve edge target_id"
def test_fan_out_edge_group_serialization(self) -> None:
"""Test that FanOutEdgeGroup can be serialized and has correct fields, including edges and type."""
edge_group = FanOutEdgeGroup(source_id="source", target_ids=["target1", "target2"])
# Test model_dump
data = edge_group.model_dump()
assert "id" in data
assert data["id"].startswith("FanOutEdgeGroup/")
# Test type field
assert "type" in data, "FanOutEdgeGroup should have 'type' field"
assert data["type"] == "FanOutEdgeGroup", f"Expected type 'FanOutEdgeGroup', got {data['type']}"
# Test selection_func_name field (should be None when no selection function is provided)
assert "selection_func_name" in data, "FanOutEdgeGroup should have 'selection_func_name' field"
assert data["selection_func_name"] is None, (
"selection_func_name should be None when no selection function is provided"
)
# Verify edges field is present and contains the correct edges
assert "edges" in data, "FanOutEdgeGroup should have 'edges' field"
assert len(data["edges"]) == 2, "FanOutEdgeGroup should have exactly two edges"
edges = data["edges"]
sources = [edge["source_id"] for edge in edges]
targets = [edge["target_id"] for edge in edges]
assert all(source == "source" for source in sources), f"All edges should have source 'source', got {sources}"
assert set(targets) == {"target1", "target2"}, f"Expected targets {{'target1', 'target2'}}, got {set(targets)}"
# Test model_dump_json
json_str = edge_group.model_dump_json()
parsed = json.loads(json_str)
assert "id" in parsed
assert parsed["id"].startswith("FanOutEdgeGroup/")
# Test type field in JSON
assert "type" in parsed, "JSON should have 'type' field"
assert parsed["type"] == "FanOutEdgeGroup", "JSON should preserve type field"
# Test selection_func_name field in JSON
assert "selection_func_name" in parsed, "JSON should have 'selection_func_name' field"
assert parsed["selection_func_name"] is None, (
"JSON selection_func_name should be None when no selection function is provided"
)
# Verify edges are preserved in JSON
assert "edges" in parsed, "JSON should have 'edges' field"
assert len(parsed["edges"]) == 2, "JSON should have exactly two edges"
json_edges = parsed["edges"]
json_sources = [edge["source_id"] for edge in json_edges]
json_targets = [edge["target_id"] for edge in json_edges]
assert all(source == "source" for source in json_sources), "JSON should preserve edge sources"
assert set(json_targets) == {"target1", "target2"}, "JSON should preserve edge targets"
def test_fan_out_edge_group_serialization_with_selection_func(self) -> None:
"""Test that FanOutEdgeGroup with named selection function serializes selection_func_name correctly."""
def custom_selector(data: Any, targets: list[str]) -> list[str]:
"""Custom selection function for testing."""
return targets[:1] # Select only the first target
edge_group = FanOutEdgeGroup(
source_id="source", target_ids=["target1", "target2"], selection_func=custom_selector
)
# Test model_dump
data = edge_group.model_dump()
assert "selection_func_name" in data, "FanOutEdgeGroup should have 'selection_func_name' field"
assert data["selection_func_name"] == "custom_selector", (
f"Expected selection_func_name 'custom_selector', got {data['selection_func_name']}"
)
# Test model_dump_json
json_str = edge_group.model_dump_json()
parsed = json.loads(json_str)
assert "selection_func_name" in parsed, "JSON should have 'selection_func_name' field"
assert parsed["selection_func_name"] == "custom_selector", "JSON should preserve selection_func_name"
def test_fan_out_edge_group_serialization_with_lambda_selection_func(self) -> None:
"""Test that FanOutEdgeGroup with lambda selection function serializes selection_func_name as '<lambda>'."""
edge_group = FanOutEdgeGroup(
source_id="source", target_ids=["target1", "target2"], selection_func=lambda data, targets: targets[:1]
)
# Test model_dump
data = edge_group.model_dump()
assert "selection_func_name" in data, "FanOutEdgeGroup should have 'selection_func_name' field"
assert data["selection_func_name"] == "<lambda>", (
f"Expected selection_func_name '<lambda>', got {data['selection_func_name']}"
)
# Test model_dump_json
json_str = edge_group.model_dump_json()
parsed = json.loads(json_str)
assert "selection_func_name" in parsed, "JSON should have 'selection_func_name' field"
assert parsed["selection_func_name"] == "<lambda>", "JSON should preserve selection_func_name as '<lambda>'"
def test_fan_in_edge_group_serialization(self) -> None:
"""Test that FanInEdgeGroup can be serialized and has correct fields, including edges and type."""
edge_group = FanInEdgeGroup(source_ids=["source1", "source2"], target_id="target")
# Test model_dump
data = edge_group.model_dump()
assert "id" in data
assert data["id"].startswith("FanInEdgeGroup/")
# Test type field
assert "type" in data, "FanInEdgeGroup should have 'type' field"
assert data["type"] == "FanInEdgeGroup", f"Expected type 'FanInEdgeGroup', got {data['type']}"
# Verify edges field is present and contains the correct edges
assert "edges" in data, "FanInEdgeGroup should have 'edges' field"
assert len(data["edges"]) == 2, "FanInEdgeGroup should have exactly two edges"
edges = data["edges"]
sources = [edge["source_id"] for edge in edges]
targets = [edge["target_id"] for edge in edges]
assert set(sources) == {"source1", "source2"}, f"Expected sources {{'source1', 'source2'}}, got {set(sources)}"
assert all(target == "target" for target in targets), f"All edges should have target 'target', got {targets}"
# Test model_dump_json
json_str = edge_group.model_dump_json()
parsed = json.loads(json_str)
assert "id" in parsed
assert parsed["id"].startswith("FanInEdgeGroup/")
# Test type field in JSON
assert "type" in parsed, "JSON should have 'type' field"
assert parsed["type"] == "FanInEdgeGroup", "JSON should preserve type field"
# Verify edges are preserved in JSON
assert "edges" in parsed, "JSON should have 'edges' field"
assert len(parsed["edges"]) == 2, "JSON should have exactly two edges"
json_edges = parsed["edges"]
json_sources = [edge["source_id"] for edge in json_edges]
json_targets = [edge["target_id"] for edge in json_edges]
assert set(json_sources) == {"source1", "source2"}, "JSON should preserve edge sources"
assert all(target == "target" for target in json_targets), "JSON should preserve edge targets"
def test_switch_case_edge_group_serialization(self) -> None:
"""Test that SwitchCaseEdgeGroup can be serialized and has correct fields, including edges and type."""
cases = [
SwitchCaseEdgeGroupCase(condition=lambda x: x > 0, target_id="positive"),
SwitchCaseEdgeGroupDefault(target_id="default"),
]
edge_group = SwitchCaseEdgeGroup(source_id="source", cases=cases)
# Test model_dump
data = edge_group.model_dump()
assert "id" in data
assert data["id"].startswith("SwitchCaseEdgeGroup/")
# Test type field
assert "type" in data, "SwitchCaseEdgeGroup should have 'type' field"
assert data["type"] == "SwitchCaseEdgeGroup", f"Expected type 'SwitchCaseEdgeGroup', got {data['type']}"
# Test cases field
assert "cases" in data, "SwitchCaseEdgeGroup should have 'cases' field"
assert len(data["cases"]) == 2, "SwitchCaseEdgeGroup should have exactly two cases"
cases_data = data["cases"]
# Check first case (SwitchCaseEdgeGroupCase)
case_obj = cases_data[0]
assert "target_id" in case_obj, "SwitchCaseEdgeGroupCase should have 'target_id' field"
assert "condition_name" in case_obj, "SwitchCaseEdgeGroupCase should have 'condition_name' field"
assert "type" in case_obj, "SwitchCaseEdgeGroupCase should have 'type' field"
assert case_obj["target_id"] == "positive", f"Expected target_id 'positive', got {case_obj['target_id']}"
assert case_obj["condition_name"] == "<lambda>", (
f"Expected condition_name '<lambda>', got {case_obj['condition_name']}"
)
assert case_obj["type"] == "Case", f"Expected type 'Case', got {case_obj['type']}"
# Check default case (SwitchCaseEdgeGroupDefault)
default_obj = cases_data[1]
assert "target_id" in default_obj, "SwitchCaseEdgeGroupDefault should have 'target_id' field"
assert "type" in default_obj, "SwitchCaseEdgeGroupDefault should have 'type' field"
assert default_obj["target_id"] == "default", f"Expected target_id 'default', got {default_obj['target_id']}"
assert default_obj["type"] == "Default", f"Expected type 'Default', got {default_obj['type']}"
# Verify edges field is present and contains the correct edges
assert "edges" in data, "SwitchCaseEdgeGroup should have 'edges' field"
assert len(data["edges"]) == 2, "SwitchCaseEdgeGroup should have exactly two edges"
edges = data["edges"]
sources = [edge["source_id"] for edge in edges]
targets = [edge["target_id"] for edge in edges]
assert all(source == "source" for source in sources), f"All edges should have source 'source', got {sources}"
assert set(targets) == {"positive", "default"}, (
f"Expected targets {{'positive', 'default'}}, got {set(targets)}"
)
# Check condition_name field in edges - SwitchCaseEdgeGroup edges don't have conditions
# because the conditional logic is implemented in the selection_func at the group level
condition_names = [edge.get("condition_name") for edge in edges]
assert all(name is None for name in condition_names), (
"SwitchCaseEdgeGroup edges should not have condition_name since conditions are handled at group level"
)
# Test model_dump_json
json_str = edge_group.model_dump_json()
parsed = json.loads(json_str)
assert "id" in parsed
assert parsed["id"].startswith("SwitchCaseEdgeGroup/")
# Test type field in JSON
assert "type" in parsed, "JSON should have 'type' field"
assert parsed["type"] == "SwitchCaseEdgeGroup", "JSON should preserve type field"
# Test cases field in JSON
assert "cases" in parsed, "JSON should have 'cases' field"
assert len(parsed["cases"]) == 2, "JSON should have exactly two cases"
json_cases = parsed["cases"]
json_case_obj = json_cases[0]
assert json_case_obj["target_id"] == "positive", "JSON should preserve case target_id"
assert json_case_obj["condition_name"] == "<lambda>", "JSON should preserve case condition_name"
assert json_case_obj["type"] == "Case", "JSON should preserve case type"
json_default_obj = json_cases[1]
assert json_default_obj["target_id"] == "default", "JSON should preserve default target_id"
assert json_default_obj["type"] == "Default", "JSON should preserve default type"
# Verify edges are preserved in JSON
assert "edges" in parsed, "JSON should have 'edges' field"
assert len(parsed["edges"]) == 2, "JSON should have exactly two edges"
json_edges = parsed["edges"]
json_sources = [edge["source_id"] for edge in json_edges]
json_targets = [edge["target_id"] for edge in json_edges]
assert all(source == "source" for source in json_sources), "JSON should preserve edge sources"
assert set(json_targets) == {"positive", "default"}, "JSON should preserve edge targets"
# Check condition_name field in JSON edges - should be None for SwitchCaseEdgeGroup
json_condition_names = [edge.get("condition_name") for edge in json_edges]
assert all(name is None for name in json_condition_names), (
"JSON SwitchCaseEdgeGroup edges should not have condition_name"
)
def test_switch_case_edge_group_serialization_with_named_condition(self) -> None:
"""Test that SwitchCaseEdgeGroup with named condition function serializes condition_name correctly."""
def is_positive(x: int) -> bool:
return x > 0
cases = [
SwitchCaseEdgeGroupCase(condition=is_positive, target_id="positive"),
SwitchCaseEdgeGroupDefault(target_id="default"),
]
edge_group = SwitchCaseEdgeGroup(source_id="source", cases=cases)
# Test model_dump
data = edge_group.model_dump()
assert "cases" in data, "SwitchCaseEdgeGroup should have 'cases' field"
cases_data = data["cases"]
case_obj = cases_data[0]
assert case_obj["condition_name"] == "is_positive", (
f"Expected condition_name 'is_positive', got {case_obj['condition_name']}"
)
# Test model_dump_json
json_str = edge_group.model_dump_json()
parsed = json.loads(json_str)
json_cases = parsed["cases"]
json_case_obj = json_cases[0]
assert json_case_obj["condition_name"] == "is_positive", "JSON should preserve named condition_name"
def test_workflow_serialization(self) -> None:
"""Test that Workflow can be serialized and has correct fields, including edges."""
executor1 = SampleExecutor(id="executor1")
executor2 = SampleExecutor(id="executor2")
workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
# Test model_dump
data = workflow.model_dump()
assert "edge_groups" in data
assert "executors" in data
assert "start_executor_id" in data
assert "max_iterations" in data
assert "workflow_id" in data
assert data["start_executor_id"] == "executor1"
assert "executor1" in data["executors"]
assert "executor2" in data["executors"]
# Verify edge groups contain edges
edge_groups = data["edge_groups"]
assert len(edge_groups) == 1, "Should have exactly one edge group"
edge_group = edge_groups[0]
assert "edges" in edge_group, "Edge group should contain 'edges' field"
assert len(edge_group["edges"]) == 1, "Should have exactly one edge"
edge = edge_group["edges"][0]
assert "source_id" in edge, "Edge should have source_id"
assert "target_id" in edge, "Edge should have target_id"
assert edge["source_id"] == "executor1", f"Expected source_id 'executor1', got {edge['source_id']}"
assert edge["target_id"] == "executor2", f"Expected target_id 'executor2', got {edge['target_id']}"
# Test model_dump_json
json_str = workflow.model_dump_json()
parsed = json.loads(json_str)
assert parsed["start_executor_id"] == "executor1"
assert "executor1" in parsed["executors"]
assert "executor2" in parsed["executors"]
# Verify edges are preserved in JSON serialization
json_edge_groups = parsed["edge_groups"]
assert len(json_edge_groups) == 1, "JSON should have exactly one edge group"
json_edge_group = json_edge_groups[0]
assert "edges" in json_edge_group, "JSON edge group should contain 'edges' field"
json_edge = json_edge_group["edges"][0]
assert json_edge["source_id"] == "executor1", "JSON should preserve edge source_id"
assert json_edge["target_id"] == "executor2", "JSON should preserve edge target_id"
def test_workflow_serialization_excludes_non_serializable_fields(self) -> None:
"""Test that non-serializable fields are excluded from serialization."""
executor1 = SampleExecutor(id="executor1")
executor2 = SampleExecutor(id="executor2")
workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
# Test model_dump - should not include private runtime objects
data = workflow.model_dump()
# These private runtime fields should not be in the serialized data
assert "_runner_context" not in data
assert "_shared_state" not in data
assert "_runner" not in data
def test_executor_field_validation(self) -> None:
"""Test that Executor field validation works correctly."""
# Valid executor
executor = SampleExecutor(id="valid-id")
assert executor.id == "valid-id"
# Test validation failure for empty id - pydantic automatically validates min_length=1
from pydantic import ValidationError
with pytest.raises(ValidationError):
SampleExecutor(id="")
def test_edge_field_validation(self) -> None:
"""Test that Edge field validation works correctly."""
# Valid edge
edge = Edge(source_id="source", target_id="target")
assert edge.source_id == "source"
assert edge.target_id == "target"
# Test validation failure for empty source_id
from pydantic import ValidationError
with pytest.raises(ValidationError):
Edge(source_id="", target_id="target")
# Test validation failure for empty target_id
with pytest.raises(ValidationError):
Edge(source_id="source", target_id="")
def test_comprehensive_edge_groups_workflow_serialization() -> None:
"""Test serialization of a workflow that uses all edge group types: SwitchCase, FanOut, and FanIn."""
from agent_framework_workflow._edge import Case, Default
# Create executors for a comprehensive workflow
router = SampleExecutor(id="router")
processor_a = SampleExecutor(id="proc_a")
processor_b = SampleExecutor(id="proc_b")
fanout_hub = SampleExecutor(id="fanout_hub")
parallel_1 = SampleExecutor(id="parallel_1")
parallel_2 = SampleExecutor(id="parallel_2")
aggregator = SampleAggregator(id="aggregator")
# Build workflow with all three edge group types
workflow = (
WorkflowBuilder()
.set_start_executor(router)
# 1. SwitchCaseEdgeGroup: Conditional routing
.add_switch_case_edge_group(
router,
[
Case(condition=lambda msg: len(str(msg)) < 10, target=processor_a),
Default(target=processor_b),
],
)
# 2. Direct edges
.add_edge(processor_a, fanout_hub)
.add_edge(processor_b, fanout_hub)
# 3. FanOutEdgeGroup: One-to-many distribution
.add_fan_out_edges(fanout_hub, [parallel_1, parallel_2])
# 4. FanInEdgeGroup: Many-to-one aggregation
.add_fan_in_edges([parallel_1, parallel_2], aggregator)
.build()
)
# Test workflow serialization
data = workflow.model_dump()
# Verify basic workflow structure
assert "edge_groups" in data
assert "executors" in data
assert "start_executor_id" in data
assert data["start_executor_id"] == "router"
# Verify all executors are present
expected_executors = {"router", "proc_a", "proc_b", "fanout_hub", "parallel_1", "parallel_2", "aggregator"}
assert set(data["executors"].keys()) == expected_executors
# Verify edge groups contain all three types
edge_groups = data["edge_groups"]
edge_group_types = [eg.get("id", "").split("/")[0] for eg in edge_groups]
# Should have: SwitchCaseEdgeGroup, SingleEdgeGroup (x2), FanOutEdgeGroup, FanInEdgeGroup
assert "SwitchCaseEdgeGroup" in edge_group_types, f"Expected SwitchCaseEdgeGroup in {edge_group_types}"
assert "FanOutEdgeGroup" in edge_group_types, f"Expected FanOutEdgeGroup in {edge_group_types}"
assert "FanInEdgeGroup" in edge_group_types, f"Expected FanInEdgeGroup in {edge_group_types}"
assert "SingleEdgeGroup" in edge_group_types, f"Expected SingleEdgeGroup in {edge_group_types}"
# Test JSON serialization
json_str = workflow.model_dump_json()
parsed = json.loads(json_str)
# Verify JSON structure matches model_dump
assert parsed["start_executor_id"] == "router"
assert set(parsed["executors"].keys()) == expected_executors
assert len(parsed["edge_groups"]) == len(edge_groups)
# Verify that serialization excludes non-serializable fields
assert "_runner_context" not in data
assert "_shared_state" not in data
assert "_runner" not in data
# Test that we can identify each edge group type by examining their structure
switch_case_groups = [eg for eg in edge_groups if eg.get("id", "").startswith("SwitchCaseEdgeGroup/")]
fan_out_groups = [eg for eg in edge_groups if eg.get("id", "").startswith("FanOutEdgeGroup/")]
fan_in_groups = [eg for eg in edge_groups if eg.get("id", "").startswith("FanInEdgeGroup/")]
single_groups = [eg for eg in edge_groups if eg.get("id", "").startswith("SingleEdgeGroup/")]
assert len(switch_case_groups) == 1, f"Expected 1 SwitchCaseEdgeGroup, got {len(switch_case_groups)}"
assert len(fan_out_groups) == 1, f"Expected 1 FanOutEdgeGroup, got {len(fan_out_groups)}"
assert len(fan_in_groups) == 1, f"Expected 1 FanInEdgeGroup, got {len(fan_in_groups)}"
assert len(single_groups) == 2, f"Expected 2 SingleEdgeGroups, got {len(single_groups)}"
# The key validation is that all edge group types are present and serializable
# Individual edge group fields may vary based on implementation,
# but each should have at least an 'id' field that identifies its type and 'edges' field
for group_type, groups in [
("SwitchCaseEdgeGroup", switch_case_groups),
("FanOutEdgeGroup", fan_out_groups),
("FanInEdgeGroup", fan_in_groups),
("SingleEdgeGroup", single_groups),
]:
for group in groups:
assert "id" in group, f"{group_type} should have 'id' field"
assert group["id"].startswith(f"{group_type}/"), f"{group_type} id should start with '{group_type}/'"
assert "edges" in group, f"{group_type} should have 'edges' field"
assert isinstance(group["edges"], list), f"{group_type} 'edges' should be a list"
assert len(group["edges"]) > 0, f"{group_type} should have at least one edge"
# Verify each edge has required fields
for edge in group["edges"]:
assert "source_id" in edge, f"{group_type} edge should have 'source_id'"
assert "target_id" in edge, f"{group_type} edge should have 'target_id'"
assert isinstance(edge["source_id"], str), f"{group_type} edge source_id should be string"
assert isinstance(edge["target_id"], str), f"{group_type} edge target_id should be string"
assert len(edge["source_id"]) > 0, f"{group_type} edge source_id should not be empty"
assert len(edge["target_id"]) > 0, f"{group_type} edge target_id should not be empty"
# Verify specific edge group edge counts
assert len(switch_case_groups[0]["edges"]) == 2, "SwitchCaseEdgeGroup should have 2 edges (proc_a and proc_b)"
assert len(fan_out_groups[0]["edges"]) == 2, "FanOutEdgeGroup should have 2 edges (parallel_1 and parallel_2)"
assert len(fan_in_groups[0]["edges"]) == 2, "FanInEdgeGroup should have 2 edges (from parallel_1 and parallel_2)"
for single_group in single_groups:
assert len(single_group["edges"]) == 1, "Each SingleEdgeGroup should have exactly 1 edge"
python/packages/workflow/tests/test_validation.py
+9 -6
View File
@@ -161,12 +161,14 @@ def test_graph_connectivity_isolated_executors():
# Create edges that include an isolated executor (self-loop that's not connected to main graph)
edge_groups = [
SingleEdgeGroup(executor1, executor2),
SingleEdgeGroup(executor3, executor3),
SingleEdgeGroup(executor1.id, executor2.id),
SingleEdgeGroup(executor3.id, executor3.id),
] # Self-loop to include in graph
executors: dict[str, Executor] = {executor1.id: executor1, executor2.id: executor2, executor3.id: executor3}
with pytest.raises(GraphConnectivityError) as exc_info:
validate_workflow_graph(edge_groups, executor1)
validate_workflow_graph(edge_groups, executors, executor1)
assert "unreachable" in str(exc_info.value).lower()
assert "executor3" in str(exc_info.value)
@@ -243,15 +245,16 @@ def test_type_compatibility_inheritance():
def test_direct_validation_function():
executor1 = StringExecutor(id="executor1")
executor2 = StringExecutor(id="executor2")
edge_groups = [SingleEdgeGroup(executor1, executor2)]
edge_groups = [SingleEdgeGroup(executor1.id, executor2.id)]
executors: dict[str, Executor] = {executor1.id: executor1, executor2.id: executor2}
# This should not raise any exceptions
validate_workflow_graph(edge_groups, executor1)
validate_workflow_graph(edge_groups, executors, executor1)
# Test with invalid start executor
executor3 = StringExecutor(id="executor3")
with pytest.raises(GraphConnectivityError):
validate_workflow_graph(edge_groups, executor3)
validate_workflow_graph(edge_groups, executors, executor3)
def test_fan_out_validation():
python/packages/workflow/tests/test_workflow.py
+163 -45
View File
@@ -22,33 +22,31 @@ from agent_framework_workflow import Message
@dataclass
class MockMessage:
class NumberMessage:
"""A mock message for testing purposes."""
data: int
class MockExecutor(Executor):
"""A mock executor for testing purposes."""
class IncrementExecutor(Executor):
"""An executor that increments message data by a specified amount for testing purposes."""
def __init__(self, id: str, limit: int = 10):
"""Initialize the mock executor with a limit."""
super().__init__(id=id)
self.limit = limit
limit: int = 10
increment: int = 1
@handler
async def mock_handler(self, message: MockMessage, ctx: WorkflowContext[MockMessage]) -> None:
async def mock_handler(self, message: NumberMessage, ctx: WorkflowContext[NumberMessage]) -> None:
if message.data < self.limit:
await ctx.send_message(MockMessage(data=message.data + 1))
await ctx.send_message(NumberMessage(data=message.data + self.increment))
else:
await ctx.add_event(WorkflowCompletedEvent(data=message.data))
class MockAggregator(Executor):
class AggregatorExecutor(Executor):
"""A mock executor that aggregates results from multiple executors."""
@handler
async def mock_handler(self, messages: list[MockMessage], ctx: WorkflowContext[Any]) -> None:
async def mock_handler(self, messages: list[NumberMessage], ctx: WorkflowContext[Any]) -> None:
# This mock simply returns the data incremented by 1
await ctx.add_event(WorkflowCompletedEvent(data=sum(msg.data for msg in messages)))
@@ -64,25 +62,25 @@ class MockExecutorRequestApproval(Executor):
"""A mock executor that simulates a request for approval."""
@handler
async def mock_handler_a(self, message: MockMessage, ctx: WorkflowContext[RequestInfoMessage]) -> None:
async def mock_handler_a(self, message: NumberMessage, ctx: WorkflowContext[RequestInfoMessage]) -> None:
"""A mock handler that requests approval."""
await ctx.set_shared_state(self.id, message.data)
await ctx.send_message(RequestInfoMessage())
@handler
async def mock_handler_b(self, message: ApprovalMessage, ctx: WorkflowContext[MockMessage]) -> None:
async def mock_handler_b(self, message: ApprovalMessage, ctx: WorkflowContext[NumberMessage]) -> None:
"""A mock handler that processes the approval response."""
data = await ctx.get_shared_state(self.id)
if message.approved:
await ctx.add_event(WorkflowCompletedEvent(data=data))
else:
await ctx.send_message(MockMessage(data=data))
await ctx.send_message(NumberMessage(data=data))
async def test_workflow_run_streaming():
"""Test the workflow run stream."""
executor_a = MockExecutor(id="executor_a")
executor_b = MockExecutor(id="executor_b")
executor_a = IncrementExecutor(id="executor_a")
executor_b = IncrementExecutor(id="executor_b")
workflow = (
WorkflowBuilder()
@@ -93,7 +91,7 @@ async def test_workflow_run_streaming():
)
result: int | None = None
async for event in workflow.run_streaming(MockMessage(data=0)):
async for event in workflow.run_streaming(NumberMessage(data=0)):
assert isinstance(event, WorkflowEvent)
if isinstance(event, WorkflowCompletedEvent):
result = event.data
@@ -103,8 +101,8 @@ async def test_workflow_run_streaming():
async def test_workflow_run_stream_not_completed():
"""Test the workflow run stream."""
executor_a = MockExecutor(id="executor_a")
executor_b = MockExecutor(id="executor_b")
executor_a = IncrementExecutor(id="executor_a")
executor_b = IncrementExecutor(id="executor_b")
workflow = (
WorkflowBuilder()
@@ -116,14 +114,14 @@ async def test_workflow_run_stream_not_completed():
)
with pytest.raises(RuntimeError):
async for _ in workflow.run_streaming(MockMessage(data=0)):
async for _ in workflow.run_streaming(NumberMessage(data=0)):
pass
async def test_workflow_run():
"""Test the workflow run."""
executor_a = MockExecutor(id="executor_a")
executor_b = MockExecutor(id="executor_b")
executor_a = IncrementExecutor(id="executor_a")
executor_b = IncrementExecutor(id="executor_b")
workflow = (
WorkflowBuilder()
@@ -133,7 +131,7 @@ async def test_workflow_run():
.build()
)
events = await workflow.run(MockMessage(data=0))
events = await workflow.run(NumberMessage(data=0))
completed_event = events.get_completed_event()
assert isinstance(completed_event, WorkflowCompletedEvent)
assert completed_event.data == 10
@@ -141,8 +139,8 @@ async def test_workflow_run():
async def test_workflow_run_not_completed():
"""Test the workflow run."""
executor_a = MockExecutor(id="executor_a")
executor_b = MockExecutor(id="executor_b")
executor_a = IncrementExecutor(id="executor_a")
executor_b = IncrementExecutor(id="executor_b")
workflow = (
WorkflowBuilder()
@@ -154,12 +152,12 @@ async def test_workflow_run_not_completed():
)
with pytest.raises(RuntimeError):
await workflow.run(MockMessage(data=0))
await workflow.run(NumberMessage(data=0))
async def test_workflow_send_responses_streaming():
"""Test the workflow run with approval."""
executor_a = MockExecutor(id="executor_a")
executor_a = IncrementExecutor(id="executor_a")
executor_b = MockExecutorRequestApproval(id="executor_b")
request_info_executor = RequestInfoExecutor()
@@ -174,7 +172,7 @@ async def test_workflow_send_responses_streaming():
)
request_info_event: RequestInfoEvent | None = None
async for event in workflow.run_streaming(MockMessage(data=0)):
async for event in workflow.run_streaming(NumberMessage(data=0)):
if isinstance(event, RequestInfoEvent):
request_info_event = event
@@ -191,7 +189,7 @@ async def test_workflow_send_responses_streaming():
async def test_workflow_send_responses():
"""Test the workflow run with approval."""
executor_a = MockExecutor(id="executor_a")
executor_a = IncrementExecutor(id="executor_a")
executor_b = MockExecutorRequestApproval(id="executor_b")
request_info_executor = RequestInfoExecutor()
@@ -205,7 +203,7 @@ async def test_workflow_send_responses():
.build()
)
events = await workflow.run(MockMessage(data=0))
events = await workflow.run(NumberMessage(data=0))
request_info_events = events.get_request_info_events()
assert len(request_info_events) == 1
@@ -219,15 +217,15 @@ async def test_workflow_send_responses():
async def test_fan_out():
"""Test a fan-out workflow."""
executor_a = MockExecutor(id="executor_a")
executor_b = MockExecutor(id="executor_b", limit=1)
executor_c = MockExecutor(id="executor_c", limit=2) # This executor will not complete the workflow
executor_a = IncrementExecutor(id="executor_a")
executor_b = IncrementExecutor(id="executor_b", limit=1)
executor_c = IncrementExecutor(id="executor_c", limit=2) # This executor will not complete the workflow
workflow = (
WorkflowBuilder().set_start_executor(executor_a).add_fan_out_edges(executor_a, [executor_b, executor_c]).build()
)
events = await workflow.run(MockMessage(data=0))
events = await workflow.run(NumberMessage(data=0))
# Each executor will emit two events: ExecutorInvokeEvent and ExecutorCompletedEvent
# executor_b will also emit a WorkflowCompletedEvent
@@ -239,15 +237,15 @@ async def test_fan_out():
async def test_fan_out_multiple_completed_events():
"""Test a fan-out workflow with multiple completed events."""
executor_a = MockExecutor(id="executor_a")
executor_b = MockExecutor(id="executor_b", limit=1)
executor_c = MockExecutor(id="executor_c", limit=1)
executor_a = IncrementExecutor(id="executor_a")
executor_b = IncrementExecutor(id="executor_b", limit=1)
executor_c = IncrementExecutor(id="executor_c", limit=1)
workflow = (
WorkflowBuilder().set_start_executor(executor_a).add_fan_out_edges(executor_a, [executor_b, executor_c]).build()
)
events = await workflow.run(MockMessage(data=0))
events = await workflow.run(NumberMessage(data=0))
# Each executor will emit two events: ExecutorInvokeEvent and ExecutorCompletedEvent
# executor_a and executor_b will also emit a WorkflowCompletedEvent
@@ -259,10 +257,10 @@ async def test_fan_out_multiple_completed_events():
async def test_fan_in():
"""Test a fan-in workflow."""
executor_a = MockExecutor(id="executor_a")
executor_b = MockExecutor(id="executor_b")
executor_c = MockExecutor(id="executor_c")
aggregator = MockAggregator(id="aggregator")
executor_a = IncrementExecutor(id="executor_a")
executor_b = IncrementExecutor(id="executor_b")
executor_c = IncrementExecutor(id="executor_c")
aggregator = AggregatorExecutor(id="aggregator")
workflow = (
WorkflowBuilder()
@@ -272,7 +270,7 @@ async def test_fan_in():
.build()
)
events = await workflow.run(MockMessage(data=0))
events = await workflow.run(NumberMessage(data=0))
# Each executor will emit two events: ExecutorInvokeEvent and ExecutorCompletedEvent
# aggregator will also emit a WorkflowCompletedEvent
@@ -289,7 +287,7 @@ def simple_executor() -> Executor:
async def handle_message(self, message: Message, context: WorkflowContext[None]) -> None:
pass
return SimpleExecutor("test_executor")
return SimpleExecutor(id="test_executor")
async def test_workflow_with_checkpointing_enabled(simple_executor: Executor):
@@ -521,7 +519,7 @@ async def test_workflow_multiple_runs_no_state_collision():
storage = FileCheckpointStorage(temp_dir)
# Create executor that tracks state in shared state
state_executor = StateTrackingExecutor("state_executor")
state_executor = StateTrackingExecutor(id="state_executor")
# Build workflow with checkpointing
workflow = (
@@ -555,3 +553,123 @@ async def test_workflow_multiple_runs_no_state_collision():
assert completed1.data != completed2.data
assert completed2.data != completed3.data
assert completed1.data != completed3.data
async def test_comprehensive_edge_groups_workflow():
"""Test a workflow that uses SwitchCaseEdgeGroup, FanOutEdgeGroup, and FanInEdgeGroup."""
from agent_framework_workflow._edge import Case, Default
# Create 6 executors for different roles with different increment values
router = IncrementExecutor(id="router", limit=1000, increment=1) # Increment by 1
processor_a = IncrementExecutor(id="proc_a", limit=1000, increment=1) # Increment by 1
processor_b = IncrementExecutor(id="proc_b", limit=1000, increment=2) # Increment by 2 (different from proc_a)
fanout_hub = IncrementExecutor(id="fanout_hub", limit=1000, increment=1) # Increment by 1
parallel_1 = IncrementExecutor(id="parallel_1", limit=1000, increment=3) # Increment by 3
parallel_2 = IncrementExecutor(
id="parallel_2", limit=1000, increment=5
) # Increment by 5 (different from parallel_1)
aggregator = AggregatorExecutor(id="aggregator") # Combines results from parallel processors
# Build workflow with different edge group types:
# 1. SwitchCase: router -> (proc_a if data < 5, else proc_b)
# 2. Direct edge: proc_a -> fanout_hub, proc_b -> fanout_hub
# 3. FanOut: fanout_hub -> [parallel_1, parallel_2]
# 4. FanIn: [parallel_1, parallel_2] -> aggregator
workflow = (
WorkflowBuilder()
.set_start_executor(router)
# Switch-case routing based on message data
.add_switch_case_edge_group(
router,
[
Case(condition=lambda msg: msg.data < 5, target=processor_a),
Default(target=processor_b),
],
)
# Both processors send to fanout hub
.add_edge(processor_a, fanout_hub)
.add_edge(processor_b, fanout_hub)
# Fan out to parallel processors
.add_fan_out_edges(fanout_hub, [parallel_1, parallel_2])
# Fan in to aggregator
.add_fan_in_edges([parallel_1, parallel_2], aggregator)
.build()
)
# Test with small number (should go through processor_a)
# router(2->3) -> switch routes to proc_a -> proc_a(3->4) -> fanout_hub(4->5)
# -> [parallel_1(5->8), parallel_2(5->10)] -> aggregator(8+10=18)
events_small = await workflow.run(NumberMessage(data=2))
completed_small = events_small.get_completed_event()
assert completed_small is not None
assert completed_small.data == 18 # Exact expected result: 8+10 from parallel processors
# Test with large number (should go through processor_b)
# router(8->9) -> switch routes to proc_b -> proc_b(9->11) -> fanout_hub(11->12)
# -> [parallel_1(12->15), parallel_2(12->17)] -> aggregator(15+17=32)
events_large = await workflow.run(NumberMessage(data=8))
completed_large = events_large.get_completed_event()
assert completed_large is not None
assert completed_large.data == 32 # Exact expected result: 15+17 from parallel processors
# The key verification is that we successfully executed a workflow using all three edge group types
# and that both switch-case paths work (small vs large numbers)
# Verify we had multiple events indicating complex execution path
assert len(events_small) >= 6 # Should have multiple executors involved
assert len(events_large) >= 6
# Verify different paths were taken by checking exact results
assert completed_small.data == 18, f"Small number path should result in 18, got {completed_small.data}"
assert completed_large.data == 32, f"Large number path should result in 32, got {completed_large.data}"
assert completed_small.data != completed_large.data, "Different paths should produce different results"
# Both tests should complete successfully, proving all edge group types work
# Additional verification: check that the workflow contains the expected edge group types
edge_groups = workflow.edge_groups
has_switch_case = any(edge_group.__class__.__name__ == "SwitchCaseEdgeGroup" for edge_group in edge_groups)
has_fan_out = any(edge_group.__class__.__name__ == "FanOutEdgeGroup" for edge_group in edge_groups)
has_fan_in = any(edge_group.__class__.__name__ == "FanInEdgeGroup" for edge_group in edge_groups)
assert has_switch_case, "Workflow should contain SwitchCaseEdgeGroup"
assert has_fan_out, "Workflow should contain FanOutEdgeGroup"
assert has_fan_in, "Workflow should contain FanInEdgeGroup"
async def test_workflow_with_simple_cycle_and_exit_condition():
"""Test a simpler workflow with a cycle that has a clear exit condition."""
# Create a simple cycle: A -> B -> A, with A having an exit condition
executor_a = IncrementExecutor(id="exec_a", limit=6, increment=2) # Exit when data >= 6
executor_b = IncrementExecutor(id="exec_b", limit=1000, increment=1) # Never exit, just increment
# Simple cycle: A -> B -> A, A exits when limit reached
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
.add_edge(executor_a, executor_b) # A -> B
.add_edge(executor_b, executor_a) # B -> A (creates cycle)
.build()
)
# Test the cycle
# Expected: exec_a(2->4) -> exec_b(4->5) -> exec_a(5->7, completes because 7 >= 6)
events = await workflow.run(NumberMessage(data=2))
completed_event = events.get_completed_event()
assert completed_event is not None
assert (
completed_event.data is not None and completed_event.data >= 6
) # Should complete when executor_a reaches its limit
# Verify cycling occurred (should have events from both executors)
# Check for ExecutorInvokeEvent and ExecutorCompletedEvent types that have executor_id
from agent_framework_workflow._events import ExecutorCompletedEvent, ExecutorInvokeEvent
executor_events = [e for e in events if isinstance(e, (ExecutorInvokeEvent, ExecutorCompletedEvent))]
executor_ids = {e.executor_id for e in executor_events}
assert "exec_a" in executor_ids, "Should have events from executor A"
assert "exec_b" in executor_ids, "Should have events from executor B"
# Should have multiple events due to cycling
assert len(events) >= 4, f"Expected at least 4 events due to cycling, got {len(events)}"
python/packages/workflow/tests/test_workflow_builder.py
+1 -1
View File
@@ -61,5 +61,5 @@ def test_workflow_builder_fluent_api():
)
assert len(workflow.edge_groups) == 4
assert workflow.start_executor.id == executor_a.id
assert workflow.start_executor_id == executor_a.id
assert len(workflow.executors) == 6