Codex
129 lines
4.6 KiB
Python
129 lines
4.6 KiB
Python
import unittest
|
|
from datetime import datetime, timezone
|
|
|
|
from app.models.knowledge_base import KnowledgeBase
|
|
from app.models.knowledge_base_document import KnowledgeBaseDocument
|
|
from app.services.knowledge_base_service import KnowledgeBaseService
|
|
|
|
|
|
class StubKnowledgeBaseRepository:
|
|
def __init__(self):
|
|
self.items = {}
|
|
|
|
def create(self, knowledge_base):
|
|
self.items[knowledge_base.id] = knowledge_base
|
|
return knowledge_base
|
|
|
|
def get_by_id(self, knowledge_base_id):
|
|
return self.items.get(knowledge_base_id)
|
|
|
|
def list_paginated(self, tenant_id, page, page_size):
|
|
return [item for item in self.items.values() if item.tenant_id == tenant_id]
|
|
|
|
def count(self, tenant_id):
|
|
return len([item for item in self.items.values() if item.tenant_id == tenant_id])
|
|
|
|
|
|
class StubKnowledgeBaseDocumentRepository:
|
|
def __init__(self):
|
|
self.items = {}
|
|
|
|
def create(self, document):
|
|
if document.created_at is None:
|
|
document.created_at = datetime.now(timezone.utc)
|
|
self.items[document.id] = document
|
|
return document
|
|
|
|
def update(self, document):
|
|
if document.created_at is None:
|
|
document.created_at = datetime.now(timezone.utc)
|
|
self.items[document.id] = document
|
|
return document
|
|
|
|
def get_by_id(self, document_id, tenant_id=None, knowledge_base_id=None):
|
|
item = self.items.get(document_id)
|
|
if item is None:
|
|
return None
|
|
if tenant_id is not None and item.tenant_id != tenant_id:
|
|
return None
|
|
if knowledge_base_id is not None and item.knowledge_base_id != knowledge_base_id:
|
|
return None
|
|
return item
|
|
|
|
def list_paginated(self, tenant_id, knowledge_base_id, page, page_size):
|
|
return [
|
|
item
|
|
for item in self.items.values()
|
|
if item.tenant_id == tenant_id and item.knowledge_base_id == knowledge_base_id
|
|
]
|
|
|
|
def count(self, tenant_id, knowledge_base_id):
|
|
return len(self.list_paginated(tenant_id, knowledge_base_id, 1, 100))
|
|
|
|
def count_by_knowledge_base_ids(self, tenant_id, knowledge_base_ids):
|
|
return {knowledge_base_id: self.count(tenant_id, knowledge_base_id) for knowledge_base_id in knowledge_base_ids}
|
|
|
|
|
|
class KnowledgeBaseServiceTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.service = KnowledgeBaseService(db=None)
|
|
self.service.repository = StubKnowledgeBaseRepository()
|
|
self.service.document_repository = StubKnowledgeBaseDocumentRepository()
|
|
self.current_user = {"id": "user-admin", "tenant_id": "tenant-demo-001", "role": "admin"}
|
|
self.knowledge_base = KnowledgeBase(
|
|
id="kb-123",
|
|
tenant_id=self.current_user["tenant_id"],
|
|
name="合同知识库",
|
|
description="测试知识库",
|
|
status="active",
|
|
)
|
|
self.service.repository.create(self.knowledge_base)
|
|
|
|
def test_failed_document_retry_without_source_explains_reupload(self):
|
|
failed = self.service.upload_document(
|
|
self.current_user,
|
|
self.knowledge_base.id,
|
|
"empty.txt",
|
|
"text/plain",
|
|
b"",
|
|
)
|
|
|
|
retried = self.service.retry_document(self.current_user, self.knowledge_base.id, failed["id"])
|
|
|
|
self.assertEqual(retried["id"], failed["id"])
|
|
self.assertEqual(retried["status"], "failed")
|
|
self.assertEqual(retried["parse_status"], "failed")
|
|
self.assertEqual(retried["index_status"], "blocked")
|
|
self.assertTrue(retried["retryable"])
|
|
self.assertIn("重新上传", retried["error_message"])
|
|
|
|
def test_failed_document_retry_with_replacement_file_indexes_same_document(self):
|
|
failed = self.service.upload_document(
|
|
self.current_user,
|
|
self.knowledge_base.id,
|
|
"empty.txt",
|
|
"text/plain",
|
|
b"",
|
|
)
|
|
|
|
retried = self.service.retry_document(
|
|
self.current_user,
|
|
self.knowledge_base.id,
|
|
failed["id"],
|
|
file_name="rules-v2.txt",
|
|
mime_type="text/plain",
|
|
content="付款、保密和违约责任制度".encode("utf-8"),
|
|
)
|
|
|
|
self.assertEqual(retried["id"], failed["id"])
|
|
self.assertEqual(retried["file_name"], "rules-v2.txt")
|
|
self.assertEqual(retried["status"], "indexed")
|
|
self.assertEqual(retried["parse_status"], "completed")
|
|
self.assertEqual(retried["index_status"], "completed")
|
|
self.assertFalse(retried["retryable"])
|
|
self.assertIn("付款", retried["content_excerpt"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|