190 lines
8.6 KiB
Python
190 lines
8.6 KiB
Python
import shutil
|
||
import tempfile
|
||
import unittest
|
||
from pathlib import Path
|
||
|
||
from sqlmodel import Session, SQLModel, create_engine, select
|
||
|
||
from app import planner
|
||
from app.agents import runner
|
||
from app.rag import dedup
|
||
from app.rag.pipeline import iter_lightweight_chunks
|
||
from app.models import (
|
||
Agent,
|
||
Campaign,
|
||
CampaignPlanVersion,
|
||
KnowledgeBase,
|
||
KnowledgeFile,
|
||
KnowledgeUsage,
|
||
Segment,
|
||
)
|
||
|
||
|
||
def parsed_segments():
|
||
return {
|
||
"status": "ok",
|
||
"segments": [
|
||
{"title": "开场建立信任", "content_tag": "开场", "time_range": "00:00-00:05",
|
||
"mindmap_node": "开场：欢迎与今天主题", "product_ref": ""},
|
||
{"title": "核心心法讲解", "content_tag": "心法", "time_range": "00:05-00:10",
|
||
"mindmap_node": "主体：核心心法", "product_ref": ""},
|
||
],
|
||
"pending_confirms": [],
|
||
"notes": [],
|
||
"provider": "test",
|
||
"model": "mock-model",
|
||
}
|
||
|
||
|
||
class PlannerVersionTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.engine = create_engine("sqlite://")
|
||
SQLModel.metadata.create_all(self.engine)
|
||
self.db = Session(self.engine)
|
||
self.tempdir = Path(tempfile.mkdtemp(prefix="yawan-planner-kb-"))
|
||
source_dir = self.tempdir / "kb"
|
||
source_dir.mkdir(parents=True, exist_ok=True)
|
||
materials = {
|
||
"开场素材.txt": (
|
||
"开场信任欢迎主题素材。开场要先建立信任，用温和欢迎语说明今天主题。"
|
||
"开场阶段必须让听众感到安全、被看见、愿意继续听下去。"
|
||
"欢迎语不能夸大承诺，也不能制造焦虑，只能围绕信任和陪伴展开。"
|
||
) * 3,
|
||
"心法素材.txt": (
|
||
"核心心法讲解素材。主体部分解释核心心法，用日常语言说明练习步骤和身体感受。"
|
||
"心法讲解要回到呼吸、放松、觉察这些可验证体验，不写素材之外的功效。"
|
||
"每个结论都必须能在核心心法素材中找到来源。"
|
||
) * 3,
|
||
}
|
||
for file_name, body in materials.items():
|
||
path = source_dir / file_name
|
||
path.write_text(body, encoding="utf-8")
|
||
# 禁用本地向量：chunk_count 为 0，Planner 必须从登记源文件临时切分并预演。
|
||
self.kb = KnowledgeBase(
|
||
name="测试知识库", index_status="disabled", file_count=len(materials),
|
||
chunk_count=0, source_dir=str(source_dir), organization_id=1,
|
||
owner_user_id=0, visibility="global")
|
||
self.agent = Agent(name="测试生产 Agent", role="production", status="active",
|
||
content_tags=["开场", "心法", "收尾"], kb_ids=[1])
|
||
self.db.add_all([self.kb, self.agent])
|
||
self.db.commit()
|
||
self.db.refresh(self.kb)
|
||
self.db.refresh(self.agent)
|
||
self.agent.kb_ids = [self.kb.id]
|
||
self.db.add(self.agent)
|
||
for file_name in materials:
|
||
path = source_dir / file_name
|
||
self.db.add(KnowledgeFile(
|
||
kb_id=self.kb.id, file_name=file_name,
|
||
content_hash=dedup.file_hash(path), revision=1))
|
||
self.campaign = Campaign(
|
||
name="计划版本测试", mindmap_text="开场 -> 心法", objective="完成测试直播",
|
||
success_criteria="只使用知识库", duration_minutes=10, status="draft",
|
||
owner_user_id=0,
|
||
)
|
||
self.db.add(self.campaign)
|
||
self.db.commit()
|
||
self.db.refresh(self.campaign)
|
||
self._original_parse = planner.parse_mindmap
|
||
planner.parse_mindmap = lambda *args, **kwargs: parsed_segments()
|
||
|
||
def tearDown(self):
|
||
planner.parse_mindmap = self._original_parse
|
||
self.db.close()
|
||
self.engine.dispose()
|
||
shutil.rmtree(self.tempdir, ignore_errors=True)
|
||
|
||
def test_dag_cycle_is_rejected(self):
|
||
with self.assertRaisesRegex(ValueError, "环"):
|
||
planner.topo_waves(["n1", "n2"], [{"from": "n1", "to": "n2"}, {"from": "n2", "to": "n1"}])
|
||
|
||
def test_planner_creates_validated_version_and_approval_materializes_segments(self):
|
||
result = planner.run_plan_job(self.db, self.campaign.id)
|
||
self.assertEqual(result["status"], "waiting_review")
|
||
self.db.refresh(self.campaign)
|
||
self.assertEqual(self.campaign.status, "plan_review")
|
||
version = self.db.get(CampaignPlanVersion, result["plan_version_id"])
|
||
self.assertEqual(version.version, 1)
|
||
self.assertEqual(len(version.plan["nodes"]), 2)
|
||
self.assertEqual(version.plan["edges"], [{"from": "n1", "to": "n2"}])
|
||
self.assertEqual(version.validation["waves"], [["n1"], ["n2"]])
|
||
self.assertTrue(version.validation["ok"])
|
||
|
||
approved = planner.approve_plan(self.db, version.id, reviewer_user_id=1)
|
||
self.db.refresh(self.campaign)
|
||
self.assertEqual(approved.status, "approved")
|
||
self.assertEqual(self.campaign.status, "parsed")
|
||
self.assertEqual(self.campaign.approved_plan_version_id, approved.id)
|
||
self.assertEqual(self.campaign.canvas_snapshot["plan_version_id"], approved.id)
|
||
segments = self.db.exec(select(Segment).where(Segment.campaign_id == self.campaign.id)).all()
|
||
self.assertEqual([s.node_id for s in segments], ["n1", "n2"])
|
||
self.assertEqual([s.agent_id for s in segments], [self.agent.id, self.agent.id])
|
||
self.assertEqual([s.content_tag for s in segments], ["开场", "心法"])
|
||
|
||
def test_worker_rejects_legacy_campaign_without_approved_plan_in_dsh_mode(self):
|
||
self.campaign.status = "parsed"
|
||
self.db.add(self.campaign)
|
||
self.db.commit()
|
||
outcome = runner.run_campaign(self.db, self.campaign.id)
|
||
self.assertEqual(outcome["status"], "error")
|
||
self.assertIn("人工批准", outcome["error"])
|
||
|
||
def test_worker_rejects_canvas_that_drifted_from_approved_version(self):
|
||
result = planner.run_plan_job(self.db, self.campaign.id)
|
||
planner.approve_plan(self.db, result["plan_version_id"], reviewer_user_id=1)
|
||
self.db.refresh(self.campaign)
|
||
self.campaign.status = "running"
|
||
snapshot = dict(self.campaign.canvas_snapshot)
|
||
snapshot["plan_version_id"] = 999999
|
||
self.campaign.canvas_snapshot = snapshot
|
||
self.db.add(self.campaign)
|
||
self.db.commit()
|
||
outcome = runner.run_campaign(self.db, self.campaign.id)
|
||
self.assertEqual(outcome["status"], "error")
|
||
self.assertIn("画布快照", outcome["error"])
|
||
|
||
def test_empty_knowledge_base_fails_validation_and_never_materializes_segments(self):
|
||
self.kb.file_count = 0
|
||
self.kb.chunk_count = 0
|
||
self.db.add(self.kb)
|
||
self.db.commit()
|
||
result = planner.run_plan_job(self.db, self.campaign.id)
|
||
self.assertEqual(result["status"], "failed")
|
||
self.db.refresh(self.campaign)
|
||
self.assertEqual(self.campaign.status, "plan_failed")
|
||
version = self.db.get(CampaignPlanVersion, result["plan_version_id"])
|
||
self.assertTrue(version.validation["errors"])
|
||
self.assertEqual(self.db.exec(select(Segment)).all(), [])
|
||
with self.assertRaisesRegex(ValueError, "不能批准"):
|
||
planner.approve_plan(self.db, version.id)
|
||
|
||
def test_consumed_lightweight_chunks_block_new_plan_before_worker_runs(self):
|
||
chunks = iter_lightweight_chunks(self.db, [self.kb.id])
|
||
self.assertEqual(len(chunks), 2)
|
||
for item in chunks:
|
||
self.db.add(KnowledgeUsage(
|
||
agent_id=self.agent.id, kb_id=self.kb.id, chunk_id=item["id"],
|
||
file_name=item["meta"].get("file", ""), excerpt=item["text"][:300],
|
||
campaign_id=999, segment_id=999, status="consumed"))
|
||
self.db.commit()
|
||
|
||
result = planner.run_plan_job(self.db, self.campaign.id)
|
||
self.assertEqual(result["status"], "failed")
|
||
version = self.db.get(CampaignPlanVersion, result["plan_version_id"])
|
||
codes = {item["code"] for item in version.validation["errors"]}
|
||
self.assertTrue({"knowledge_exhausted", "knowledge_mismatch"} & codes)
|
||
self.assertEqual(self.db.exec(select(Segment)).all(), [])
|
||
|
||
def test_unassigned_agent_is_blocking_error(self):
|
||
self.agent.status = "disabled"
|
||
self.db.add(self.agent)
|
||
self.db.commit()
|
||
result = planner.run_plan_job(self.db, self.campaign.id)
|
||
version = self.db.get(CampaignPlanVersion, result["plan_version_id"])
|
||
codes = {item["code"] for item in version.validation["errors"]}
|
||
self.assertIn("unassigned_agent", codes)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|