138 lines
8.2 KiB
Python
138 lines
8.2 KiB
Python
|
import unittest
|
||
|
|
from pathlib import Path
|
||
|
|
from tempfile import TemporaryDirectory
|
||
|
|
from unittest.mock import patch
|
||
|
|
|
||
|
|
from fastapi import HTTPException
|
||
|
|
from sqlmodel import SQLModel, Session, create_engine
|
||
|
|
|
||
|
|
from app.agents.snapshot import agent_from_snapshot, freeze_agent_snapshot, resolve_snapshot_provider
|
||
|
|
from app.auth import (authenticate_token, has_permission, hash_password, issue_session,
|
||
|
|
required_permission, revoke_token, user_payload, verify_password)
|
||
|
|
from app.main import create_campaign, get_campaign, get_trace, list_traces
|
||
|
|
from app.models import (Agent, AgentPromptTemplate, AppUser, Campaign, CampaignRun, LLMCall, ModelProvider,
|
||
|
|
KnowledgeBase, RunEvent, AgentSession, RoleDefinition)
|
||
|
|
from app.rag import pipeline
|
||
|
|
|
||
|
|
|
||
|
|
class GovernanceTests(unittest.TestCase):
|
||
|
|
def setUp(self):
|
||
|
|
self.engine = create_engine("sqlite://")
|
||
|
|
SQLModel.metadata.create_all(self.engine)
|
||
|
|
self.db = Session(self.engine)
|
||
|
|
|
||
|
|
def tearDown(self):
|
||
|
|
self.db.close()
|
||
|
|
|
||
|
|
def test_password_session_and_rbac_contract(self):
|
||
|
|
password_hash = hash_password("correct-horse-123")
|
||
|
|
self.assertTrue(verify_password("correct-horse-123", password_hash))
|
||
|
|
self.assertFalse(verify_password("wrong-password", password_hash))
|
||
|
|
user = AppUser(username="admin", display_name="管理员", password_hash=password_hash, role="admin")
|
||
|
|
self.db.add(user); self.db.commit(); self.db.refresh(user)
|
||
|
|
token, _ = issue_session(self.db, user)
|
||
|
|
self.assertEqual(authenticate_token(self.db, f"Bearer {token}").id, user.id)
|
||
|
|
self.assertTrue(has_permission(user, "users:manage"))
|
||
|
|
self.assertEqual(required_permission("POST", "/api/auth/logout"), "read")
|
||
|
|
self.assertEqual(required_permission("POST", "/api/kbs/3/rebuild"), "kb:write")
|
||
|
|
self.assertEqual(required_permission("GET", "/api/tags"), "read")
|
||
|
|
self.assertEqual(required_permission("POST", "/api/kbs"), "kb:write")
|
||
|
|
self.assertEqual(required_permission("DELETE", "/api/memories/3"), "memories:write")
|
||
|
|
self.assertEqual(user_payload(user)["portal"], "admin")
|
||
|
|
revoke_token(self.db, f"Bearer {token}")
|
||
|
|
self.assertIsNone(authenticate_token(self.db, f"Bearer {token}"))
|
||
|
|
|
||
|
|
def test_campaign_business_contract_persists_and_is_bounded(self):
|
||
|
|
created = create_campaign({
|
||
|
|
"name": "任务契约测试",
|
||
|
|
"objective": "完成一段可直接口播的心法内容",
|
||
|
|
"success_criteria": "只使用知识库；引用可核对；不扩展第二个主题",
|
||
|
|
"mindmap_text": "开场 -> 心法 -> 收尾",
|
||
|
|
}, s=self.db, request=None)
|
||
|
|
detail = get_campaign(created["id"], s=self.db, request=None)
|
||
|
|
self.assertEqual(detail["objective"], "完成一段可直接口播的心法内容")
|
||
|
|
self.assertEqual(detail["success_criteria"], "只使用知识库；引用可核对；不扩展第二个主题")
|
||
|
|
self.assertEqual(self.db.get(Campaign, created["id"]).mindmap_text, "开场 -> 心法 -> 收尾")
|
||
|
|
with self.assertRaises(HTTPException) as context:
|
||
|
|
create_campaign({"objective": "目" * 501, "success_criteria": "可验收"},
|
||
|
|
s=self.db, request=None)
|
||
|
|
self.assertEqual(context.exception.status_code, 422)
|
||
|
|
# 验收标准允许 AI/默认稿自动生成：缺失时不再 422，而是落一份系统默认契约。
|
||
|
|
auto = create_campaign({"objective": "只有目标"}, s=self.db, request=None)
|
||
|
|
detail_auto = get_campaign(auto["id"], s=self.db, request=None)
|
||
|
|
self.assertTrue(detail_auto["success_criteria"].strip())
|
||
|
|
|
||
|
|
def test_custom_role_with_console_permission_uses_admin_portal(self):
|
||
|
|
user = AppUser(username="custom-admin", display_name="内容平台主管",
|
||
|
|
password_hash=hash_password("correct-horse-123"), role="custom_a1b2c3d4e5f6")
|
||
|
|
self.db.add(user); self.db.commit(); self.db.refresh(user)
|
||
|
|
self.db.add(RoleDefinition(role=user.role, label="内容平台主管",
|
||
|
|
permissions=["read", "roles:manage", "audit:manage"], organization_id=1))
|
||
|
|
self.db.commit()
|
||
|
|
payload = user_payload(user, self.db)
|
||
|
|
self.assertEqual(payload["portal"], "admin")
|
||
|
|
self.assertIn("roles:manage", payload["permissions"])
|
||
|
|
|
||
|
|
def test_snapshot_blocks_provider_drift(self):
|
||
|
|
provider = ModelProvider(name="测试 Provider", base_url="https://gateway.example/v1",
|
||
|
|
models=["model-a"], default_model="model-a", enabled=True)
|
||
|
|
self.db.add(provider); self.db.commit(); self.db.refresh(provider)
|
||
|
|
agent = Agent(name="冻结质检", role="qc", rules="旧规则", model_provider_id=provider.id,
|
||
|
|
model_name="model-a")
|
||
|
|
template = AgentPromptTemplate(name="冻结模板", role_prompt="你是{agent_name}",
|
||
|
|
task_template="标题：{title}", version=4,
|
||
|
|
output_schema={"type": "object", "required": ["result"]})
|
||
|
|
self.db.add(template); self.db.commit(); self.db.refresh(template)
|
||
|
|
agent.prompt_template_id = template.id
|
||
|
|
self.db.add(agent); self.db.commit(); self.db.refresh(agent)
|
||
|
|
snapshot = freeze_agent_snapshot(self.db, agent)
|
||
|
|
agent.rules = "新规则"
|
||
|
|
self.assertEqual(agent_from_snapshot(snapshot).rules, "旧规则")
|
||
|
|
self.assertEqual(snapshot["prompt_template"]["version"], 4)
|
||
|
|
self.assertEqual(snapshot["prompt_template"]["task_template"], "标题：{title}")
|
||
|
|
self.assertEqual(resolve_snapshot_provider(self.db, snapshot)[1], "model-a")
|
||
|
|
provider.base_url = "https://changed.example/v1"; self.db.add(provider); self.db.commit()
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
resolve_snapshot_provider(self.db, snapshot)
|
||
|
|
|
||
|
|
def test_trace_aggregates_calls_sessions_and_events(self):
|
||
|
|
run = CampaignRun(campaign_id=8, trace_id="trace-governance", status="awaiting_review")
|
||
|
|
call = LLMCall(campaign_id=8, trace_id="trace-governance", agent_name="质检", provider="test",
|
||
|
|
model="model-a", status="ok", prompt_tokens=12, completion_tokens=8, duration_ms=55)
|
||
|
|
session = AgentSession(campaign_id=8, trace_id="trace-governance", agent_name="质检", phase="qc_style")
|
||
|
|
event = RunEvent(run_id=1, campaign_id=8, trace_id="trace-governance", event_type="model_finished", status="ok")
|
||
|
|
self.db.add(run); self.db.commit(); self.db.refresh(run)
|
||
|
|
event.run_id = run.id or 0
|
||
|
|
self.db.add(call); self.db.add(session); self.db.add(event); self.db.commit()
|
||
|
|
rows = list_traces(s=self.db)["items"]
|
||
|
|
self.assertEqual(rows[0]["trace_id"], "trace-governance")
|
||
|
|
self.assertEqual(rows[0]["call_count"], 1)
|
||
|
|
detail = get_trace("trace-governance", s=self.db)
|
||
|
|
self.assertEqual(detail["run"]["id"], run.id)
|
||
|
|
self.assertEqual(len(detail["calls"]), 1)
|
||
|
|
self.assertEqual(len(detail["sessions"]), 1)
|
||
|
|
|
||
|
|
def test_rebuild_failure_restores_previous_vector_snapshot(self):
|
||
|
|
with TemporaryDirectory() as temp:
|
||
|
|
source = Path(temp)
|
||
|
|
(source / "source.md").write_text("可恢复素材" * 20, encoding="utf-8")
|
||
|
|
kb = KnowledgeBase(name="回滚测试库", source_dir=str(source), index_status="ready",
|
||
|
|
file_count=1, chunk_count=7)
|
||
|
|
self.db.add(kb); self.db.commit(); self.db.refresh(kb)
|
||
|
|
snapshot = {"ids": ["chunk-old"], "documents": ["旧素材"],
|
||
|
|
"embeddings": [[0.1, 0.2]], "metadatas": [{"file": "source.md"}]}
|
||
|
|
with patch.object(pipeline.store, "snapshot_kb", return_value=snapshot), \
|
||
|
|
patch.object(pipeline.store, "reset_kb"), \
|
||
|
|
patch.object(pipeline.store, "restore_kb") as restore, \
|
||
|
|
patch.object(pipeline.embed, "embed_texts", side_effect=RuntimeError("embedding failed")):
|
||
|
|
with self.assertRaises(RuntimeError):
|
||
|
|
pipeline.rebuild_kb(self.db, kb.id or 0)
|
||
|
|
restore.assert_called_once_with(kb.id, snapshot)
|
||
|
|
self.db.refresh(kb)
|
||
|
|
self.assertEqual(kb.index_status, "ready")
|
||
|
|
self.assertEqual(kb.chunk_count, 7)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|