import unittest from fastapi import HTTPException from fastapi.testclient import TestClient from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker from sqlalchemy.pool import StaticPool from app.api.knowledge_bases import get_audit_log_service, get_shared_ai_platform_client, get_tenant_service, require_current_user from app.db import Base, get_db_session from app.main import app class StubAuditLogService: def __init__(self): self.events = [] def record_event(self, **kwargs): self.events.append(kwargs) return kwargs class StubKnowledgeBaseService: def list_knowledge_bases(self, current_user, page, page_size): if current_user["role"] != "admin": raise HTTPException(status_code=403, detail="forbidden") return { "items": [ { "id": "kb-123", "name": "劳动合同知识库", "description": "劳动用工条款", "status": "active", "document_count": 1, "retrieval_mode": "hybrid", "created_at": "2026-04-16T00:00:00+00:00", "updated_at": "2026-04-16T00:00:00+00:00", } ], "page": page, "page_size": page_size, "total": 1, } def create_knowledge_base(self, current_user, name, description): if current_user["role"] != "admin": raise HTTPException(status_code=403, detail="forbidden") return { "id": "kb-123", "name": name, "description": description, "status": "active", "document_count": 0, "retrieval_mode": "hybrid", "created_at": "2026-04-16T00:00:00+00:00", "updated_at": "2026-04-16T00:00:00+00:00", } def list_knowledge_base_versions(self, current_user, knowledge_base_id): if current_user["role"] != "admin": raise HTTPException(status_code=403, detail="forbidden") return [ { "id": "kbver-001", "knowledge_base_id": knowledge_base_id, "version_label": "v1", "name": "劳动合同知识库", "description": "劳动用工条款", "retrieval_mode": "hybrid", "status": "active", "published_at": "2026-04-16T00:00:00+00:00", "published_by": "user-admin", "change_summary": "publish", "created_at": "2026-04-16T00:00:00+00:00", } ] def publish_knowledge_base(self, current_user, knowledge_base_id): if current_user["role"] != "admin": raise HTTPException(status_code=403, detail="forbidden") return { "id": knowledge_base_id, "name": "劳动合同知识库", "description": "劳动用工条款", "status": "active", "document_count": 1, "retrieval_mode": "hybrid", "current_version_label": "v1", "published_at": "2026-05-15T00:00:00+00:00", "published_by": current_user["id"], "created_at": "2026-04-16T00:00:00+00:00", "updated_at": "2026-05-15T00:00:00+00:00", } def unpublish_knowledge_base(self, current_user, knowledge_base_id): if current_user["role"] != "admin": raise HTTPException(status_code=403, detail="forbidden") result = self.publish_knowledge_base(current_user, knowledge_base_id) result["status"] = "archived" result["published_at"] = None result["published_by"] = None return result def rollback_knowledge_base(self, current_user, knowledge_base_id, target_version_label, change_summary=None): if current_user["role"] != "admin": raise HTTPException(status_code=403, detail="forbidden") result = self.publish_knowledge_base(current_user, knowledge_base_id) result["current_version_label"] = target_version_label return result def list_documents(self, current_user, knowledge_base_id, page, page_size): if current_user["role"] != "admin": raise HTTPException(status_code=403, detail="forbidden") return { "items": [ { "id": "kbdoc-123", "knowledge_base_id": knowledge_base_id, "file_name": "rules.txt", "mime_type": "text/plain", "file_size": 12, "status": "indexed", "source_type": "statute", "authority": "全国人民代表大会", "effective_date": "2021-01-01", "region": "CN", "tags": ["合同法", "通用"], "content_excerpt": "合同解除规则", "extracted_text_length": 6, "preview_truncated": False, "created_at": "2026-04-16T00:00:00+00:00", "current_stage": "indexed", "current_stage_label": "已完成索引", "status_detail": "文本提取与索引已完成，可直接用于检索和问答。", "parse_status": "completed", "index_status": "completed", "failure_reason": None, "next_action": "可直接进入问答，或继续补充更多制度与条款文档。", "retryable": False, }, { "id": "kbdoc-456", "knowledge_base_id": knowledge_base_id, "file_name": "draft.docx", "mime_type": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", "file_size": 18, "status": "processing", "source_type": "policy", "authority": "企业法务部", "effective_date": "2026-05-01", "region": "CN-SH", "tags": ["制度"], "content_excerpt": None, "extracted_text_length": 0, "preview_truncated": False, "created_at": "2026-04-16T00:10:00+00:00", "current_stage": "parse", "current_stage_label": "文本提取中", "status_detail": "系统正在提取文本并写入检索索引。", "parse_status": "processing", "index_status": "blocked", "failure_reason": "文本提取尚未完成，索引阶段暂时阻塞。", "next_action": "请等待文本提取完成；提取成功后系统会继续进入索引阶段。", "retryable": False, } ], "page": page, "page_size": page_size, "total": 2, } def upload_document(self, current_user, knowledge_base_id, file_name, mime_type, content): if current_user["role"] != "admin": raise HTTPException(status_code=403, detail="forbidden") return { "id": "kbdoc-123", "knowledge_base_id": knowledge_base_id, "file_name": file_name, "mime_type": mime_type, "file_size": len(content), "status": "indexed", "content_excerpt": content.decode("utf-8"), "extracted_text_length": len(content.decode("utf-8")), "preview_truncated": False, "created_at": "2026-04-16T00:00:00+00:00", "current_stage": "indexed", "current_stage_label": "已完成索引", "status_detail": "文本提取与索引已完成，可直接用于检索和问答。", "parse_status": "completed", "index_status": "completed", "failure_reason": None, "next_action": "可直接进入问答，或继续补充更多制度与条款文档。", "error_message": None, } def retry_document(self, current_user, knowledge_base_id, document_id, file_name=None, mime_type=None, content=None): if current_user["role"] != "admin": raise HTTPException(status_code=403, detail="forbidden") indexed = content is not None and len(content) > 0 return { "id": document_id, "knowledge_base_id": knowledge_base_id, "file_name": file_name or "draft.docx", "mime_type": mime_type or "text/plain", "file_size": len(content or b""), "status": "indexed" if indexed else "failed", "content_excerpt": content.decode("utf-8") if indexed else None, "extracted_text_length": len(content.decode("utf-8")) if indexed else 0, "preview_truncated": False, "created_at": "2026-04-16T00:00:00+00:00", "current_stage": "indexed" if indexed else "failed", "current_stage_label": "已完成索引" if indexed else "文本提取失败", "status_detail": "文本提取与索引已完成，可直接用于检索和问答。" if indexed else "重试需要重新上传可复制文本版本，当前记录未保留可重新索引的原文。", "parse_status": "completed" if indexed else "failed", "index_status": "completed" if indexed else "blocked", "failure_reason": None if indexed else "重试需要重新上传可复制文本版本，当前记录未保留可重新索引的原文。", "next_action": "可直接进入问答，或继续补充更多制度与条款文档。" if indexed else "请检查文件格式、内容编码或改为可复制文本后重新上传。", "retryable": not indexed, "error_message": None if indexed else "重试需要重新上传可复制文本版本，当前记录未保留可重新索引的原文。", } class StubTenantService: def __init__(self, blocked_action=None): self.blocked_action = blocked_action self.calls = [] def assert_billing_allows(self, current_user, action, requested_file_size_bytes=0): self.calls.append((action, requested_file_size_bytes)) if action == self.blocked_action: raise HTTPException(status_code=402, detail="知识库或存储额度已用完，请升级套餐或购买超额包后继续使用。") class KnowledgeBasesApiTestCase(unittest.TestCase): def setUp(self): self.audit_log_service = StubAuditLogService() self.tenant_service = StubTenantService() self.engine = create_engine( "sqlite://", future=True, connect_args={"check_same_thread": False}, poolclass=StaticPool, ) self.testing_session_local = sessionmaker(autocommit=False, autoflush=False, bind=self.engine) Base.metadata.create_all(bind=self.engine) def override_db(): db = self.testing_session_local() try: yield db finally: db.close() app.dependency_overrides[get_db_session] = override_db app.dependency_overrides[get_shared_ai_platform_client] = lambda: StubKnowledgeBaseService() app.dependency_overrides[get_audit_log_service] = lambda: self.audit_log_service app.dependency_overrides[get_tenant_service] = lambda: self.tenant_service self.client = TestClient(app) def tearDown(self): app.dependency_overrides.clear() Base.metadata.drop_all(bind=self.engine) self.engine.dispose() def test_admin_can_list_knowledge_bases(self): app.dependency_overrides[require_current_user] = lambda: { "id": "user-admin", "tenant_id": "tenant-demo-001", "role": "admin", } response = self.client.get("/api/v1/knowledge-bases") self.assertEqual(response.status_code, 200) self.assertEqual(response.json()["data"]["items"][0]["name"], "劳动合同知识库") def test_admin_can_list_documents_with_blocked_index_stage(self): app.dependency_overrides[require_current_user] = lambda: { "id": "user-admin", "tenant_id": "tenant-demo-001", "role": "admin", } response = self.client.get("/api/v1/knowledge-bases/kb-123/documents") self.assertEqual(response.status_code, 200) items = response.json()["data"]["items"] self.assertEqual(len(items), 2) blocked_item = items[1] self.assertEqual(blocked_item["current_stage"], "parse") self.assertEqual(blocked_item["index_status"], "blocked") self.assertEqual(blocked_item["failure_reason"], "文本提取尚未完成，索引阶段暂时阻塞。") self.assertIn("继续进入索引阶段", blocked_item["next_action"]) self.assertEqual(items[0]["authority"], "全国人民代表大会") self.assertEqual(items[0]["region"], "CN") def test_admin_can_create_knowledge_base(self): app.dependency_overrides[require_current_user] = lambda: { "id": "user-admin", "tenant_id": "tenant-demo-001", "role": "admin", } response = self.client.post( "/api/v1/knowledge-bases", data={"name": "采购合同知识库", "description": "采购模板与制度"}, ) self.assertEqual(response.status_code, 200) body = response.json() self.assertEqual(body["data"]["name"], "采购合同知识库") self.assertEqual(self.tenant_service.calls[0][0], "knowledge_base") self.assertEqual(self.audit_log_service.events[0]["action"], "create_knowledge_base") self.assertEqual( self.audit_log_service.events[0]["detail_json"], { "name": "采购合同知识库", "description": "采购模板与制度", "status": "active", "document_count": 0, }, ) def test_admin_can_upload_document(self): app.dependency_overrides[require_current_user] = lambda: { "id": "user-admin", "tenant_id": "tenant-demo-001", "role": "admin", } response = self.client.post( "/api/v1/knowledge-bases/kb-123/documents", data={"note": "首批规则"}, files={"file": ("rules.txt", "合同解除规则".encode("utf-8"), "text/plain")}, ) self.assertEqual(response.status_code, 200) body = response.json() self.assertEqual(body["data"]["file_name"], "rules.txt") self.assertEqual(self.tenant_service.calls[0][0], "storage") self.assertEqual(self.tenant_service.calls[0][1], len("合同解除规则".encode("utf-8"))) self.assertEqual(self.audit_log_service.events[0]["action"], "upload_knowledge_document") self.assertEqual( self.audit_log_service.events[0]["detail_json"], { "knowledge_base_id": "kb-123", "file_name": "rules.txt", "mime_type": "text/plain", "file_size": len("合同解除规则".encode("utf-8")), "status": "indexed", "current_stage": "indexed", "parse_status": "completed", "index_status": "completed", "error_message": None, "note": "首批规则", "source_type": None, "authority": None, "effective_date": None, "region": None, "tags": None, }, ) def test_create_knowledge_base_returns_billing_block_message(self): app.dependency_overrides[require_current_user] = lambda: { "id": "user-admin", "tenant_id": "tenant-demo-001", "role": "admin", } app.dependency_overrides[get_tenant_service] = lambda: StubTenantService(blocked_action="knowledge_base") response = self.client.post( "/api/v1/knowledge-bases", data={"name": "采购合同知识库", "description": "采购模板与制度"}, ) self.assertEqual(response.status_code, 402) self.assertIn("额度已用完", response.json()["message"]) def test_admin_can_publish_unpublish_rollback_and_list_versions(self): app.dependency_overrides[require_current_user] = lambda: { "id": "user-admin", "tenant_id": "tenant-demo-001", "role": "admin", } versions_response = self.client.get("/api/v1/knowledge-bases/kb-123/versions") publish_response = self.client.post("/api/v1/knowledge-bases/kb-123/publish") unpublish_response = self.client.post("/api/v1/knowledge-bases/kb-123/unpublish") rollback_response = self.client.post( "/api/v1/knowledge-bases/kb-123/rollback", data={"target_version_label": "v1", "change_summary": "rollback"}, ) self.assertEqual(versions_response.status_code, 200) self.assertEqual(versions_response.json()["data"][0]["version_label"], "v1") self.assertEqual(publish_response.status_code, 200) self.assertEqual(publish_response.json()["data"]["published_by"], "user-admin") self.assertEqual(unpublish_response.status_code, 200) self.assertEqual(unpublish_response.json()["data"]["status"], "archived") self.assertEqual(rollback_response.status_code, 200) self.assertEqual(rollback_response.json()["data"]["current_version_label"], "v1") self.assertEqual(self.audit_log_service.events[-3]["action"], "publish_knowledge_base") self.assertEqual(self.audit_log_service.events[-2]["action"], "unpublish_knowledge_base") self.assertEqual(self.audit_log_service.events[-1]["action"], "rollback_knowledge_base") def test_admin_can_retry_document_with_replacement_file(self): app.dependency_overrides[require_current_user] = lambda: { "id": "user-admin", "tenant_id": "tenant-demo-001", "role": "admin", } response = self.client.post( "/api/v1/knowledge-bases/kb-123/documents/kbdoc-456/retry", data={"note": "替换为空文档"}, files={"file": ("rules-v2.txt", "补充合同审查规则".encode("utf-8"), "text/plain")}, ) self.assertEqual(response.status_code, 200) body = response.json() self.assertEqual(body["data"]["id"], "kbdoc-456") self.assertEqual(body["data"]["status"], "indexed") self.assertEqual(self.audit_log_service.events[0]["action"], "retry_knowledge_document") self.assertEqual( self.audit_log_service.events[0]["detail_json"], { "knowledge_base_id": "kb-123", "original_document_id": "kbdoc-456", "file_name": "rules-v2.txt", "mime_type": "text/plain", "file_size": len("补充合同审查规则".encode("utf-8")), "status": "indexed", "current_stage": "indexed", "parse_status": "completed", "index_status": "completed", "error_message": None, "replacement_file": True, "note": "替换为空文档", "source_type": None, "authority": None, "effective_date": None, "region": None, "tags": None, }, ) def test_uploaded_metadata_is_persisted_and_list_overlay_returns_it(self): app.dependency_overrides[require_current_user] = lambda: { "id": "user-admin", "tenant_id": "tenant-demo-001", "role": "admin", } upload_response = self.client.post( "/api/v1/knowledge-bases/kb-123/documents", data={ "note": "带法源元数据", "source_type": "法律", "authority": "全国人民代表大会", "effective_date": "2021-01-01", "region": "CN", "tags_json": '["民法典","合同"]', }, files={"file": ("rules.txt", "合同解除规则".encode("utf-8"), "text/plain")}, ) self.assertEqual(upload_response.status_code, 200) listed = self.client.get("/api/v1/knowledge-bases/kb-123/documents") self.assertEqual(listed.status_code, 200) item = listed.json()["data"]["items"][0] self.assertEqual(item["source_type"], "法律") self.assertEqual(item["authority"], "全国人民代表大会") self.assertEqual(item["effective_date"], "2021-01-01") self.assertEqual(item["region"], "CN") self.assertEqual(item["tags"], ["民法典", "合同"]) def test_invalid_tags_json_is_rejected(self): app.dependency_overrides[require_current_user] = lambda: { "id": "user-admin", "tenant_id": "tenant-demo-001", "role": "admin", } response = self.client.post( "/api/v1/knowledge-bases/kb-123/documents", data={"tags_json": '{"bad":"shape"}'}, files={"file": ("rules.txt", "abc".encode("utf-8"), "text/plain")}, ) self.assertEqual(response.status_code, 400) self.assertEqual(response.json()["message"], "tags_json must be a JSON array of strings") def test_member_cannot_access_knowledge_bases(self): app.dependency_overrides[require_current_user] = lambda: { "id": "user-member", "tenant_id": "tenant-demo-001", "role": "member", } response = self.client.get("/api/v1/knowledge-bases") self.assertEqual(response.status_code, 403) self.assertEqual(response.json()["message"], "forbidden") self.assertEqual(self.audit_log_service.events, []) if __name__ == "__main__": unittest.main()