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