289 lines
10 KiB
Python
289 lines
10 KiB
Python
"""TeamPlan / SubTask 数据模型单元测试 (hub-and-spoke 模式)"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from agentkit.experts.plan import (
|
||
MergeStrategy,
|
||
PlanStatus,
|
||
SubTask,
|
||
SubTaskStatus,
|
||
TeamPlan,
|
||
)
|
||
|
||
|
||
# ── 辅助函数 ──────────────────────────────────────────────
|
||
|
||
|
||
def _make_subtask(
|
||
id: str = "subtask_1",
|
||
description: str = "分析数据",
|
||
assigned_expert: str = "analyst",
|
||
status: SubTaskStatus = SubTaskStatus.PENDING,
|
||
result: dict | None = None,
|
||
) -> SubTask:
|
||
"""创建测试用 SubTask 实例"""
|
||
return SubTask(
|
||
id=id,
|
||
description=description,
|
||
assigned_expert=assigned_expert,
|
||
status=status,
|
||
result=result,
|
||
)
|
||
|
||
|
||
def _make_valid_plan() -> TeamPlan:
|
||
"""创建一个有效的 hub-and-spoke 执行计划"""
|
||
subtasks = [
|
||
_make_subtask(id="s1", description="分析需求", assigned_expert="analyst"),
|
||
_make_subtask(id="s2", description="设计架构", assigned_expert="architect"),
|
||
_make_subtask(id="s3", description="编写代码", assigned_expert="coder"),
|
||
]
|
||
return TeamPlan(
|
||
id="plan_001",
|
||
task="实现用户登录功能",
|
||
subtasks=subtasks,
|
||
status=PlanStatus.DRAFT,
|
||
lead_expert="architect",
|
||
)
|
||
|
||
|
||
# ── MergeStrategy 测试 ────────────────────────────────────
|
||
|
||
|
||
class TestMergeStrategy:
|
||
"""MergeStrategy 枚举测试"""
|
||
|
||
def test_only_best_strategy_exists(self):
|
||
"""hub-and-spoke 模式仅保留 BEST 策略"""
|
||
assert MergeStrategy.BEST == "best"
|
||
# 确保只有 BEST 一个值
|
||
assert len(list(MergeStrategy)) == 1
|
||
|
||
def test_no_vote_or_fusion(self):
|
||
"""VOTE 和 FUSION 已被移除"""
|
||
assert not hasattr(MergeStrategy, "VOTE")
|
||
assert not hasattr(MergeStrategy, "FUSION")
|
||
|
||
|
||
# ── PlanStatus 测试 ───────────────────────────────────────
|
||
|
||
|
||
class TestPlanStatus:
|
||
"""PlanStatus 枚举测试"""
|
||
|
||
def test_statuses_exist(self):
|
||
"""必要的计划状态都存在"""
|
||
assert PlanStatus.DRAFT == "draft"
|
||
assert PlanStatus.EXECUTING == "executing"
|
||
assert PlanStatus.COMPLETED == "completed"
|
||
assert PlanStatus.FAILED == "failed"
|
||
assert PlanStatus.FALLBACK == "fallback"
|
||
|
||
def test_no_confirmed_status(self):
|
||
"""CONFIRMED 状态已被移除（hub-and-spoke 无需确认阶段）"""
|
||
assert not hasattr(PlanStatus, "CONFIRMED")
|
||
|
||
|
||
# ── SubTaskStatus 测试 ────────────────────────────────────
|
||
|
||
|
||
class TestSubTaskStatus:
|
||
"""SubTaskStatus 枚举测试"""
|
||
|
||
def test_statuses_exist(self):
|
||
"""子任务状态都存在"""
|
||
assert SubTaskStatus.PENDING == "pending"
|
||
assert SubTaskStatus.RUNNING == "running"
|
||
assert SubTaskStatus.COMPLETED == "completed"
|
||
assert SubTaskStatus.FAILED == "failed"
|
||
|
||
|
||
# ── SubTask 测试 ──────────────────────────────────────────
|
||
|
||
|
||
class TestSubTask:
|
||
"""SubTask 数据模型测试"""
|
||
|
||
def test_creation_with_all_fields(self):
|
||
"""创建 SubTask 并设置所有字段"""
|
||
subtask = SubTask(
|
||
id="subtask_a",
|
||
description="竞品分析",
|
||
assigned_expert="analyst",
|
||
status=SubTaskStatus.RUNNING,
|
||
result={"report": "竞品分析报告"},
|
||
)
|
||
assert subtask.id == "subtask_a"
|
||
assert subtask.description == "竞品分析"
|
||
assert subtask.assigned_expert == "analyst"
|
||
assert subtask.status == SubTaskStatus.RUNNING
|
||
assert subtask.result == {"report": "竞品分析报告"}
|
||
|
||
def test_default_values(self):
|
||
"""默认值：自动生成 id，PENDING 状态"""
|
||
subtask = SubTask(description="测试任务")
|
||
assert subtask.id is not None
|
||
assert len(subtask.id) > 0
|
||
assert subtask.description == "测试任务"
|
||
assert subtask.assigned_expert == ""
|
||
assert subtask.status == SubTaskStatus.PENDING
|
||
assert subtask.result is None
|
||
|
||
def test_to_dict_from_dict_roundtrip(self):
|
||
"""to_dict / from_dict 往返序列化"""
|
||
subtask = SubTask(
|
||
id="roundtrip_subtask",
|
||
description="往返测试",
|
||
assigned_expert="tester",
|
||
status=SubTaskStatus.COMPLETED,
|
||
result={"key": "value"},
|
||
)
|
||
d = subtask.to_dict()
|
||
restored = SubTask.from_dict(d)
|
||
|
||
assert restored.id == subtask.id
|
||
assert restored.description == subtask.description
|
||
assert restored.assigned_expert == subtask.assigned_expert
|
||
assert restored.status == subtask.status
|
||
assert restored.result == subtask.result
|
||
|
||
def test_to_dict_structure(self):
|
||
"""to_dict 返回正确的字典结构"""
|
||
subtask = _make_subtask(
|
||
id="struct_test",
|
||
description="结构测试",
|
||
assigned_expert="dev",
|
||
status=SubTaskStatus.RUNNING,
|
||
result={"output": "data"},
|
||
)
|
||
d = subtask.to_dict()
|
||
assert d["id"] == "struct_test"
|
||
assert d["description"] == "结构测试"
|
||
assert d["assigned_expert"] == "dev"
|
||
assert d["status"] == "running"
|
||
assert d["result"] == {"output": "data"}
|
||
|
||
|
||
# ── TeamPlan 测试 ─────────────────────────────────────────
|
||
|
||
|
||
class TestTeamPlan:
|
||
"""TeamPlan 数据模型测试"""
|
||
|
||
def test_creation(self):
|
||
"""创建 TeamPlan"""
|
||
plan = _make_valid_plan()
|
||
assert plan.id == "plan_001"
|
||
assert plan.task == "实现用户登录功能"
|
||
assert len(plan.subtasks) == 3
|
||
assert plan.status == PlanStatus.DRAFT
|
||
assert plan.lead_expert == "architect"
|
||
|
||
def test_default_values(self):
|
||
"""默认值：自动生成 id，空子任务列表"""
|
||
plan = TeamPlan(task="测试任务")
|
||
assert plan.id is not None
|
||
assert plan.task == "测试任务"
|
||
assert plan.subtasks == []
|
||
assert plan.status == PlanStatus.DRAFT
|
||
assert plan.lead_expert == ""
|
||
|
||
def test_to_dict_from_dict_roundtrip(self):
|
||
"""to_dict / from_dict 往返序列化"""
|
||
plan = _make_valid_plan()
|
||
d = plan.to_dict()
|
||
restored = TeamPlan.from_dict(d)
|
||
|
||
assert restored.id == plan.id
|
||
assert restored.task == plan.task
|
||
assert len(restored.subtasks) == len(plan.subtasks)
|
||
assert restored.status == plan.status
|
||
assert restored.lead_expert == plan.lead_expert
|
||
|
||
for original, restored_st in zip(plan.subtasks, restored.subtasks):
|
||
assert restored_st.id == original.id
|
||
assert restored_st.description == original.description
|
||
assert restored_st.assigned_expert == original.assigned_expert
|
||
assert restored_st.status == original.status
|
||
|
||
def test_get_subtask_by_id(self):
|
||
"""get_subtask 根据 ID 获取子任务"""
|
||
plan = _make_valid_plan()
|
||
st = plan.get_subtask("s2")
|
||
assert st is not None
|
||
assert st.id == "s2"
|
||
assert st.description == "设计架构"
|
||
|
||
def test_get_subtask_with_nonexistent_id_returns_none(self):
|
||
"""get_subtask 对不存在的 ID 返回 None"""
|
||
plan = _make_valid_plan()
|
||
assert plan.get_subtask("nonexistent") is None
|
||
|
||
def test_update_subtask_status(self):
|
||
"""update_subtask_status 更新子任务状态和结果"""
|
||
plan = _make_valid_plan()
|
||
plan.update_subtask_status("s1", SubTaskStatus.COMPLETED, {"output": "分析完成"})
|
||
st = plan.get_subtask("s1")
|
||
assert st is not None
|
||
assert st.status == SubTaskStatus.COMPLETED
|
||
assert st.result == {"output": "分析完成"}
|
||
|
||
def test_update_subtask_status_without_result(self):
|
||
"""update_subtask_status 不传 result 时不覆盖已有 result"""
|
||
plan = _make_valid_plan()
|
||
plan.update_subtask_status("s1", SubTaskStatus.COMPLETED, {"output": "done"})
|
||
plan.update_subtask_status("s1", SubTaskStatus.RUNNING)
|
||
st = plan.get_subtask("s1")
|
||
assert st is not None
|
||
assert st.status == SubTaskStatus.RUNNING
|
||
assert st.result == {"output": "done"}
|
||
|
||
def test_completed_subtasks_property(self):
|
||
"""completed_subtasks 返回已完成的子任务"""
|
||
plan = _make_valid_plan()
|
||
plan.update_subtask_status("s1", SubTaskStatus.COMPLETED)
|
||
plan.update_subtask_status("s2", SubTaskStatus.COMPLETED)
|
||
plan.update_subtask_status("s3", SubTaskStatus.FAILED)
|
||
|
||
completed = plan.completed_subtasks
|
||
assert len(completed) == 2
|
||
assert {st.id for st in completed} == {"s1", "s2"}
|
||
|
||
def test_failed_subtasks_property(self):
|
||
"""failed_subtasks 返回失败的子任务"""
|
||
plan = _make_valid_plan()
|
||
plan.update_subtask_status("s1", SubTaskStatus.COMPLETED)
|
||
plan.update_subtask_status("s2", SubTaskStatus.FAILED)
|
||
plan.update_subtask_status("s3", SubTaskStatus.FAILED)
|
||
|
||
failed = plan.failed_subtasks
|
||
assert len(failed) == 2
|
||
assert {st.id for st in failed} == {"s2", "s3"}
|
||
|
||
def test_all_done_property_when_all_completed(self):
|
||
"""all_done 当所有子任务完成时返回 True"""
|
||
plan = _make_valid_plan()
|
||
for st in plan.subtasks:
|
||
plan.update_subtask_status(st.id, SubTaskStatus.COMPLETED)
|
||
assert plan.all_done is True
|
||
|
||
def test_all_done_property_when_some_failed(self):
|
||
"""all_done 当所有子任务完成或失败时返回 True"""
|
||
plan = _make_valid_plan()
|
||
plan.update_subtask_status("s1", SubTaskStatus.COMPLETED)
|
||
plan.update_subtask_status("s2", SubTaskStatus.FAILED)
|
||
plan.update_subtask_status("s3", SubTaskStatus.COMPLETED)
|
||
assert plan.all_done is True
|
||
|
||
def test_all_done_property_when_pending(self):
|
||
"""all_done 当有子任务未完成时返回 False"""
|
||
plan = _make_valid_plan()
|
||
plan.update_subtask_status("s1", SubTaskStatus.COMPLETED)
|
||
# s2 and s3 still pending
|
||
assert plan.all_done is False
|
||
|
||
def test_all_done_property_empty_plan(self):
|
||
"""all_done 当没有子任务时返回 True（vacuous truth）"""
|
||
plan = TeamPlan(task="空计划")
|
||
assert plan.all_done is True
|