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()