394 lines
14 KiB
Python
394 lines
14 KiB
Python
"""
|
||
Phase 3 验证测试：重复上传旧切片残留修复
|
||
|
||
测试场景：
|
||
1. add_file_to_kb() 中的"先删后加"逻辑
|
||
2. upload_document() 中的同名文件覆盖逻辑
|
||
3. 批量上传中的同名文件处理
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import shutil
|
||
import uuid
|
||
from pathlib import Path
|
||
|
||
# 添加项目根目录到 Python 路径
|
||
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
TEST_TEMP_ROOT = os.path.join(PROJECT_ROOT, ".data", "test-runs")
|
||
sys.path.insert(0, PROJECT_ROOT)
|
||
|
||
|
||
def _make_test_dir():
|
||
"""在工作区内创建测试目录，兼容 Codex 的 Windows 文件权限。"""
|
||
os.makedirs(TEST_TEMP_ROOT, exist_ok=True)
|
||
path = os.path.join(TEST_TEMP_ROOT, f"upload-{uuid.uuid4().hex}")
|
||
os.makedirs(path)
|
||
return path
|
||
|
||
|
||
def _cleanup_test_dir(path):
|
||
root = os.path.realpath(TEST_TEMP_ROOT)
|
||
target = os.path.realpath(path)
|
||
if os.path.commonpath([root, target]) != root:
|
||
raise RuntimeError(f"拒绝清理测试目录之外的路径: {target}")
|
||
shutil.rmtree(target, ignore_errors=True)
|
||
|
||
|
||
def test_add_file_to_kb_dedup():
|
||
"""测试 add_file_to_kb() 中的去重逻辑"""
|
||
print("\n=== 测试 1：add_file_to_kb 去重逻辑 ===")
|
||
|
||
# 模拟 ChromaDB collection
|
||
class MockCollection:
|
||
def __init__(self):
|
||
self.data = {} # id -> {doc, meta, embedding}
|
||
|
||
def add(self, ids, documents, metadatas, embeddings):
|
||
for i, doc, meta, emb in zip(ids, documents, metadatas, embeddings):
|
||
self.data[i] = {'doc': doc, 'meta': meta, 'embedding': emb}
|
||
|
||
def get(self, where=None, ids=None, include=None):
|
||
results = {'ids': [], 'documents': [], 'metadatas': [], 'embeddings': []}
|
||
|
||
if ids:
|
||
for id in ids:
|
||
if id in self.data:
|
||
results['ids'].append(id)
|
||
results['documents'].append(self.data[id]['doc'])
|
||
results['metadatas'].append(self.data[id]['meta'])
|
||
results['embeddings'].append(self.data[id]['embedding'])
|
||
return results
|
||
|
||
for id, item in self.data.items():
|
||
if where:
|
||
# 简单的 where 过滤
|
||
match = True
|
||
for k, v in where.items():
|
||
if item['meta'].get(k) != v:
|
||
match = False
|
||
break
|
||
if not match:
|
||
continue
|
||
|
||
results['ids'].append(id)
|
||
results['documents'].append(item['doc'])
|
||
results['metadatas'].append(item['meta'])
|
||
results['embeddings'].append(item['embedding'])
|
||
|
||
return results
|
||
|
||
def delete(self, ids):
|
||
for id in ids:
|
||
self.data.pop(id, None)
|
||
|
||
def count(self):
|
||
return len(self.data)
|
||
|
||
# 模拟知识库管理器（仅测试去重逻辑）
|
||
class MockKBManager:
|
||
def __init__(self):
|
||
self.collections = {'test_kb': MockCollection()}
|
||
self._bm25_indexes = {}
|
||
|
||
def get_collection(self, kb_name):
|
||
return self.collections.get(kb_name)
|
||
|
||
def rebuild_bm25_index(self, kb_name):
|
||
pass
|
||
|
||
def add_file_to_kb_logic(self, kb_name, filename, chunks):
|
||
"""模拟 add_file_to_kb 的核心逻辑"""
|
||
collection = self.get_collection(kb_name)
|
||
if not collection:
|
||
return 0
|
||
|
||
# 入库前清理同名旧切片（核心修复）
|
||
existing = collection.get(where={"source": filename})
|
||
if existing and existing['ids']:
|
||
old_count = len(existing['ids'])
|
||
collection.delete(ids=existing['ids'])
|
||
print(f" [清理] 删除旧切片: {filename}, 共 {old_count} 个")
|
||
|
||
# 添加新切片
|
||
ids = []
|
||
documents = []
|
||
metadatas = []
|
||
embeddings = []
|
||
|
||
for i, chunk in enumerate(chunks):
|
||
chunk_id = f"{filename}_{i}"
|
||
ids.append(chunk_id)
|
||
documents.append(chunk)
|
||
metadatas.append({
|
||
"source": filename,
|
||
"chunk_index": i,
|
||
"collection": kb_name
|
||
})
|
||
embeddings.append([0.1] * 768) # 模拟向量
|
||
|
||
collection.add(ids, documents, metadatas, embeddings)
|
||
return len(ids)
|
||
|
||
# 测试场景 1：首次添加
|
||
kb = MockKBManager()
|
||
count1 = kb.add_file_to_kb_logic('test_kb', '规章制度.pdf', ['内容1', '内容2', '内容3'])
|
||
assert count1 == 3, f"首次添加应返回 3，实际 {count1}"
|
||
assert kb.collections['test_kb'].count() == 3, f"向量库应有 3 个切片，实际 {kb.collections['test_kb'].count()}"
|
||
print(f" [PASS] 首次添加: {count1} 个切片")
|
||
|
||
# 测试场景 2：重复上传同名文件（应清理旧切片）
|
||
count2 = kb.add_file_to_kb_logic('test_kb', '规章制度.pdf', ['新内容1', '新内容2'])
|
||
assert count2 == 2, f"重复上传应返回 2，实际 {count2}"
|
||
assert kb.collections['test_kb'].count() == 2, f"向量库应有 2 个切片（旧切片已清理），实际 {kb.collections['test_kb'].count()}"
|
||
print(f" [PASS] 重复上传: {count2} 个切片（旧切片已清理）")
|
||
|
||
# 验证切片内容是新的
|
||
result = kb.collections['test_kb'].get()
|
||
assert '新内容1' in result['documents'], "切片内容应为新内容"
|
||
assert '内容1' not in result['documents'], "旧内容应已被删除"
|
||
print(f" [PASS] 切片内容已更新为新版本")
|
||
|
||
# 测试场景 3：不同名文件不互相影响
|
||
count3 = kb.add_file_to_kb_logic('test_kb', '操作手册.pdf', ['手册1'])
|
||
assert count3 == 1, f"添加新文件应返回 1，实际 {count3}"
|
||
assert kb.collections['test_kb'].count() == 3, f"向量库应有 3 个切片（2+1），实际 {kb.collections['test_kb'].count()}"
|
||
print(f" [PASS] 不同名文件独立存在: 总计 {kb.collections['test_kb'].count()} 个切片")
|
||
|
||
|
||
def test_upload_document_overwrite():
|
||
"""测试 upload_document() 中的同名文件覆盖逻辑"""
|
||
print("\n=== 测试 2：upload_document 同名文件覆盖 ===")
|
||
|
||
# 创建临时目录
|
||
temp_dir = _make_test_dir()
|
||
try:
|
||
target_dir = os.path.join(temp_dir, 'public_kb')
|
||
os.makedirs(target_dir, exist_ok=True)
|
||
|
||
# 创建旧文件
|
||
old_file = os.path.join(target_dir, '规章制度.pdf')
|
||
with open(old_file, 'w', encoding='utf-8') as f:
|
||
f.write('旧内容')
|
||
|
||
assert os.path.exists(old_file), "旧文件应存在"
|
||
old_size = os.path.getsize(old_file)
|
||
print(f" [准备] 创建旧文件: {old_file}, 大小 {old_size} 字节")
|
||
|
||
# 模拟上传逻辑：同名文件覆盖
|
||
filename = '规章制度.pdf'
|
||
filepath = os.path.join(target_dir, filename)
|
||
|
||
replaced = False
|
||
if os.path.exists(filepath):
|
||
replaced = True
|
||
# 这里应该调用 kb_manager 清理旧切片，简化测试只验证文件覆盖
|
||
print(f" [检测] 发现同名文件，准备覆盖")
|
||
|
||
# 保存新文件（覆盖）
|
||
with open(filepath, 'w', encoding='utf-8') as f:
|
||
f.write('新内容，更长一些')
|
||
|
||
assert replaced, "应检测到同名文件并标记 replaced=True"
|
||
assert os.path.exists(filepath), "新文件应存在"
|
||
new_size = os.path.getsize(filepath)
|
||
assert new_size > old_size, f"新文件应更大（{new_size} > {old_size}）"
|
||
|
||
with open(filepath, 'r', encoding='utf-8') as f:
|
||
content = f.read()
|
||
assert content == '新内容，更长一些', "文件内容应为新内容"
|
||
|
||
print(f" [PASS] 文件已覆盖: {old_size} -> {new_size} 字节")
|
||
print(f" [PASS] replaced={replaced}")
|
||
|
||
finally:
|
||
_cleanup_test_dir(temp_dir)
|
||
|
||
|
||
def test_batch_upload_overwrite():
|
||
"""测试批量上传中的同名文件覆盖逻辑"""
|
||
print("\n=== 测试 3：批量上传同名文件覆盖 ===")
|
||
|
||
temp_dir = _make_test_dir()
|
||
try:
|
||
target_dir = os.path.join(temp_dir, 'public_kb')
|
||
os.makedirs(target_dir, exist_ok=True)
|
||
|
||
# 创建旧文件
|
||
old_file = os.path.join(target_dir, '规章制度.pdf')
|
||
with open(old_file, 'w', encoding='utf-8') as f:
|
||
f.write('旧内容')
|
||
|
||
# 模拟批量上传逻辑
|
||
files = [
|
||
('规章制度.pdf', '新内容1'),
|
||
('操作手册.pdf', '手册内容'),
|
||
]
|
||
|
||
results = []
|
||
for filename, content in files:
|
||
filepath = os.path.join(target_dir, filename)
|
||
|
||
replaced = False
|
||
if os.path.exists(filepath):
|
||
replaced = True
|
||
print(f" [检测] {filename} 已存在，准备覆盖")
|
||
|
||
with open(filepath, 'w', encoding='utf-8') as f:
|
||
f.write(content)
|
||
|
||
results.append({
|
||
'filename': filename,
|
||
'replaced': replaced
|
||
})
|
||
|
||
# 验证结果
|
||
assert len(results) == 2, f"应处理 2 个文件，实际 {len(results)}"
|
||
assert results[0]['replaced'] is True, "规章制度.pdf 应标记为 replaced"
|
||
assert results[1]['replaced'] is False, "操作手册.pdf 不应标记为 replaced"
|
||
|
||
print(f" [PASS] 批量上传: {len(results)} 个文件")
|
||
print(f" [PASS] 规章制度.pdf: replaced={results[0]['replaced']}")
|
||
print(f" [PASS] 操作手册.pdf: replaced={results[1]['replaced']}")
|
||
|
||
# 验证文件内容
|
||
with open(old_file, 'r', encoding='utf-8') as f:
|
||
content = f.read()
|
||
assert content == '新内容1', "规章制度.pdf 应为新内容"
|
||
print(f" [PASS] 文件内容已更新")
|
||
|
||
finally:
|
||
_cleanup_test_dir(temp_dir)
|
||
|
||
|
||
def test_docstore_cleanup():
|
||
"""测试 DocStore 文件清理逻辑"""
|
||
print("\n=== 测试 4：DocStore 文件清理 ===")
|
||
|
||
temp_dir = _make_test_dir()
|
||
try:
|
||
docstore_dir = Path(temp_dir) / 'docstore'
|
||
docstore_dir.mkdir()
|
||
|
||
# 创建模拟的 DocStore 文件
|
||
collection = 'public_kb'
|
||
filename = '规章制度.pdf'
|
||
|
||
# 旧切片对应的 DocStore 文件
|
||
old_files = [
|
||
f'{collection}_{filename}_0.json',
|
||
f'{collection}_{filename}_1.json',
|
||
f'{collection}_{filename}_2.json',
|
||
]
|
||
|
||
# 其他文件的 DocStore（不应被清理）
|
||
other_files = [
|
||
f'{collection}_操作手册.pdf_0.json',
|
||
f'dept_1_kb_{filename}_0.json', # 不同 collection
|
||
]
|
||
|
||
for f in old_files + other_files:
|
||
(docstore_dir / f).write_text('{"test": "data"}', encoding='utf-8')
|
||
|
||
assert len(list(docstore_dir.glob('*.json'))) == 5, "应有 5 个 DocStore 文件"
|
||
print(f" [准备] 创建 5 个 DocStore 文件")
|
||
|
||
# 模拟清理逻辑
|
||
cleaned = 0
|
||
for ds_file in docstore_dir.glob(f'{collection}_{filename}_*.json'):
|
||
ds_file.unlink()
|
||
cleaned += 1
|
||
|
||
assert cleaned == 3, f"应清理 3 个旧 DocStore 文件，实际 {cleaned}"
|
||
remaining = list(docstore_dir.glob('*.json'))
|
||
assert len(remaining) == 2, f"应剩余 2 个 DocStore 文件，实际 {len(remaining)}"
|
||
|
||
print(f" [PASS] 清理了 {cleaned} 个旧 DocStore 文件")
|
||
print(f" [PASS] 剩余 {len(remaining)} 个无关文件未被清理")
|
||
|
||
finally:
|
||
_cleanup_test_dir(temp_dir)
|
||
|
||
|
||
def test_sync_hash_cleanup():
|
||
"""测试同步哈希记录清理逻辑"""
|
||
print("\n=== 测试 5：同步哈希记录清理 ===")
|
||
|
||
# 模拟 SyncDatabase
|
||
class MockSyncDatabase:
|
||
def __init__(self):
|
||
self.hashes = {}
|
||
|
||
def set_document_hash(self, doc_id, doc_name, hash_val, size, mtime):
|
||
self.hashes[doc_id] = {
|
||
'document_id': doc_id,
|
||
'document_name': doc_name,
|
||
'content_hash': hash_val,
|
||
'file_size': size,
|
||
'last_modified': mtime
|
||
}
|
||
|
||
def get_document_hash(self, doc_id):
|
||
return self.hashes.get(doc_id)
|
||
|
||
def delete_document_hash(self, doc_id):
|
||
self.hashes.pop(doc_id, None)
|
||
|
||
db = MockSyncDatabase()
|
||
|
||
# 添加旧哈希记录
|
||
db.set_document_hash('public_kb/规章制度.pdf', '规章制度.pdf', 'old_hash_abc', 1024, '2026-01-01')
|
||
assert db.get_document_hash('public_kb/规章制度.pdf') is not None, "旧哈希应存在"
|
||
print(f" [准备] 创建旧哈希记录: public_kb/规章制度.pdf")
|
||
|
||
# 模拟上传时清理哈希
|
||
collection = 'public_kb'
|
||
filename = '规章制度.pdf'
|
||
db.delete_document_hash(f"{collection}/{filename}")
|
||
|
||
assert db.get_document_hash('public_kb/规章制度.pdf') is None, "旧哈希应已被清理"
|
||
print(f" [PASS] 哈希记录已清理")
|
||
|
||
|
||
def run_all_tests():
|
||
"""运行所有测试"""
|
||
print("=" * 50)
|
||
print("Phase 3 验证测试：重复上传旧切片残留修复")
|
||
print("=" * 50)
|
||
|
||
tests = [
|
||
test_add_file_to_kb_dedup,
|
||
test_upload_document_overwrite,
|
||
test_batch_upload_overwrite,
|
||
test_docstore_cleanup,
|
||
test_sync_hash_cleanup,
|
||
]
|
||
|
||
passed = 0
|
||
failed = 0
|
||
|
||
for test in tests:
|
||
try:
|
||
test()
|
||
passed += 1
|
||
except Exception as e:
|
||
print(f" [FAIL] {test.__name__}: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
failed += 1
|
||
|
||
print("\n" + "=" * 50)
|
||
print(f"测试结果：{passed} 通过 / {failed} 失败 / {len(tests)} 总计")
|
||
if failed == 0:
|
||
print("全部通过！")
|
||
else:
|
||
print(f"有 {failed} 个测试失败")
|
||
print("=" * 50)
|
||
|
||
return failed == 0
|
||
|
||
|
||
if __name__ == '__main__':
|
||
success = run_all_tests()
|
||
sys.exit(0 if success else 1)
|