import unittest from fastapi import HTTPException from fastapi.testclient import TestClient from app.api.knowledge_bases import get_audit_log_service, get_shared_ai_platform_client, require_current_user 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, "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, "created_at": "2026-04-16T00:00:00+00:00", "updated_at": "2026-04-16T00:00:00+00:00", } 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", "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", "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 KnowledgeBasesApiTestCase(unittest.TestCase): def setUp(self): self.audit_log_service = StubAuditLogService() app.dependency_overrides[get_shared_ai_platform_client] = lambda: StubKnowledgeBaseService() app.dependency_overrides[get_audit_log_service] = lambda: self.audit_log_service self.client = TestClient(app) def tearDown(self): app.dependency_overrides.clear() 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"]) 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.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.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": "首批规则", }, ) 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": "替换为空文档", }, ) 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()