import copy import json from pathlib import Path import tempfile import unittest from unittest.mock import patch from src.step2_knowledge_ingest.ingest import CONTRACT, DIM, KINDS, canonical, content_parts, index_text, load, route, sha from src.step2_knowledge_ingest.portable import export, import_archive, verify_archive class FakeClient: def __init__(self, data): self.data = data self.database = None def query(self, collection, **_): return [{'count(*)': len(self.data[self.database][collection])}] def upsert(self, collection, data, **_): table = self.data[self.database].setdefault(collection, {}) for row in data: table[row['id']] = copy.deepcopy(row) def flush(self, *_args, **_kwargs): pass def close(self): pass class FakeStore: data = None def __init__(self, settings): self.settings = settings self.timeout = 1 self.client = FakeClient(self.data) def use(self, kind, create): name = self.settings['MILVUS_DB_PREFIX'] + '_' + kind if create: self.data.setdefault(name, {}) elif name not in self.data: raise ValueError('missing_database') self.client.database = name return name def ensure(self, name, content=False): self.data[self.client.database].setdefault(name, {}) def get(self, name, ids): table = self.data[self.client.database][name] return [copy.deepcopy(table[i]) for i in ids if i in table] def verify(self, records, content_name, index_name): for record in records: for part in content_parts(record): self.assert_row(content_name, part) if index_text(record) is not None: self.assert_row(index_name, {'id': record['id'], 'record_hash': sha(canonical(record))}) def assert_row(self, name, expected): row = self.data[self.client.database][name][expected['id']] for key, value in expected.items(): if row[key] != value: raise ValueError('fake_verify_mismatch') def make_record(kind, ordinal): record = {'id': 'k_' + format(ordinal, '024x'), 'kind': kind, 'title': '合成标题', 'data': {'question': '合成问题', 'body': '完整正文' + '字' * 13000}, 'source': {'url': 'https://example.invalid/' + str(ordinal)}} if kind == 'qa': record.update(summary_status='not_required_qa', retrieval_text='合成问题') else: record.update(summary='合成摘要', summary_status='generated', retrieval_text='合成标题\n合成摘要') return record class PortableTests(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.addCleanup(self.tmp.cleanup) self.root = Path(self.tmp.name) self.prepared = self.root / 'prepared' self.prepared.mkdir() records = [] for i, kind in enumerate(('policy', 'finance', 'general', 'qa'), 1): record = make_record(kind, i) records.append(record) source = {k: v for k, v in record.items() if k not in ('summary', 'summary_status', 'retrieval_text')} (self.prepared / f'{kind}.jsonl').write_text(canonical(source) + '\n', encoding='utf-8') (self.prepared / f'{kind}.with_summary.jsonl').write_text(canonical(record) + '\n', encoding='utf-8') model_hash = sha(canonical({})) snapshot = sha(canonical({'contract': CONTRACT, 'model': model_hash, 'records': records})) self.manifest = {'complete': True, 'snapshot': snapshot, 'contract': CONTRACT, 'model_hash': model_hash, 'model_files': {}, 'dimension': DIM, 'counts': {route(r): 1 for r in records}, 'databases': {}} self.source_data = {} for record in records: kind = route(record) content, index = 'content_' + snapshot[:20], 'index_' + snapshot[:20] database = 'demo_' + kind self.manifest['databases'][kind] = {'database': database, 'content_collection': content, 'index_collection': index} self.source_data[database] = {content: {p['id']: p for p in content_parts(record)}, index: {record['id']: {'id': record['id'], 'vector': [1.0] + [0.0] * (DIM - 1), 'index_text': index_text(record), 'record_hash': sha(canonical(record))}}} ingest_dir = self.prepared / 'ingestion' ingest_dir.mkdir() (ingest_dir / 'manifest.json').write_text(canonical(self.manifest), encoding='utf-8') self.settings = {'MILVUS_DB_PREFIX': 'demo', 'EMBEDDING_MODEL_PATH': self.root / 'model'} self.archive = self.root / 'snapshot.zip' def test_export_verify_and_restore(self): FakeStore.data = self.source_data snapshot, count = export(self.prepared, self.archive, self.settings, FakeStore) self.assertEqual((snapshot, count), (self.manifest['snapshot'], 4)) self.assertEqual(verify_archive(self.archive)[1], 4) target = self.root / 'target' FakeStore.data = {} with patch('src.step2_knowledge_ingest.portable.model_matches_manifest', return_value=True): restored, count = import_archive(self.archive, target, self.settings, FakeStore) self.assertEqual((restored, count), (snapshot, 4)) self.assertEqual(sum(len(db[next(k for k in db if k.startswith('index_'))]) for db in FakeStore.data.values()), 4) self.assertEqual(json.loads((target / 'ingestion/manifest.json').read_text(encoding='utf-8'))['snapshot'], snapshot) def test_missing_remote_part_rejects_and_leaves_no_archive(self): FakeStore.data = self.source_data first = next(iter(self.source_data.values())) content = next(k for k in first if k.startswith('content_')) first[content].pop(next(iter(first[content]))) with self.assertRaises(ValueError): export(self.prepared, self.archive, self.settings, FakeStore) self.assertFalse(self.archive.exists()) self.assertFalse((self.root / 'snapshot.zip.partial').exists()) def test_corrupt_archive_rejected_before_target_write(self): FakeStore.data = self.source_data export(self.prepared, self.archive, self.settings, FakeStore) from zipfile import ZipFile with ZipFile(self.archive, 'a') as archive: archive.writestr('untrusted.txt', 'bad') with self.assertRaises(ValueError): verify_archive(self.archive) FakeStore.data = {} with patch('src.step2_knowledge_ingest.portable.model_matches_manifest', return_value=True): with self.assertRaises(ValueError): import_archive(self.archive, self.root / 'target', self.settings, FakeStore) self.assertEqual(FakeStore.data, {}) def version(self, records, name, database=None, model_files=None): """Produce independently versioned synthetic packages with real hashing.""" prepared = self.root / (name + '-source') prepared.mkdir() ordered = [] for kind in KINDS: subset = [r for r in records if r['kind'] == kind] ordered.extend(subset) raw = [{k: v for k, v in r.items() if k not in ('summary', 'summary_status', 'retrieval_text')} for r in subset] for suffix, rows in (('', raw), ('.with_summary', subset)): (prepared / f'{kind}{suffix}.jsonl').write_text(''.join(canonical(r) + '\n' for r in rows), encoding='utf-8') from collections import Counter model_hash = sha(canonical(model_files or {})) snapshot = sha(canonical({'contract': CONTRACT, 'model': model_hash, 'records': ordered})) manifest = {**self.manifest, 'model_hash': model_hash, 'model_files': model_files or {}, 'snapshot': snapshot, 'counts': dict(Counter(route(r) for r in records)), 'databases': {}} data = {} for record in records: kind = route(record) content, index = 'content_' + snapshot[:20], 'index_' + snapshot[:20] db = 'demo_' + kind manifest['databases'][kind] = {'database': db, 'content_collection': content, 'index_collection': index} tables = data.setdefault(db, {content: {}, index: {}}) tables[content].update({p['id']: p for p in content_parts(record)}) tables[index][record['id']] = {'id':record['id'], 'vector':[1.0] + [0.0] * (DIM-1), 'index_text':index_text(record), 'record_hash':sha(canonical(record))} (prepared/'ingestion').mkdir() (prepared/'ingestion/manifest.json').write_text(canonical(manifest), encoding='utf-8') package = self.root / (name + '.zip') FakeStore.data = data export(prepared, package, self.settings, FakeStore, database=database) return package def six_records(self): rows = load(self.prepared) guide = make_record('general', 5) guide['collection'] = '办事指南' service = make_record('general', 6) service['collection'] = '企呼我应六大服务板块' return rows + [guide, service] def restore_package(self, package, target, data, database=None, store=FakeStore): FakeStore.data = data with patch('src.step2_knowledge_ingest.portable.model_matches_manifest', return_value=True): return import_archive(package, target, self.settings, store, database=database) def test_full_update_different_version_and_reimport_preserves_backups(self): rows = self.six_records() first = self.version(rows, 'first') second = self.version(rows + [make_record('general', 9)], 'second') target, data = self.root/'target', {} self.restore_package(first, target, data) old = (target/'ingestion/manifest.json').read_bytes() old_collections = {db:set(tables) for db,tables in data.items()} new_snapshot, count = self.restore_package(second, target, data) self.assertEqual(count, 7) self.assertEqual(len(load(target)), 7) self.assertTrue(any((p/'ingestion/manifest.json').read_bytes() == old for p in self.root.glob('target.previous-*'))) for db, names in old_collections.items(): self.assertTrue(names <= data[db].keys()) self.assertEqual(self.restore_package(second, target, data), (new_snapshot, 7)) self.assertEqual(len(load(target)), 7) def test_single_library_replaces_removed_rows_and_preserves_other_five(self): rows = self.six_records() + [make_record('general', 9)] first = self.version(rows, 'first') replacement = make_record('general', 10) second = self.version([replacement], 'article-only', database='article') target, data = self.root/'target', {} self.restore_package(first, target, data) before = copy.deepcopy(data) previous = json.loads((target/'ingestion/manifest.json').read_text()) snapshot, imported = self.restore_package(second, target, data, 'article') current = json.loads((target/'ingestion/manifest.json').read_text()) actual = load(target) self.assertEqual(imported, 1) self.assertEqual([r for r in actual if route(r)=='article'], [replacement]) self.assertEqual(len(actual), 6) for kind in previous['databases']: if kind != 'article': self.assertEqual(data['demo_'+kind], before['demo_'+kind]) self.assertEqual(current['databases'][kind]['content_collection'], previous['databases'][kind]['content_collection']) # A mixed live snapshot can be exported to a standalone full package, # then imported on a fresh server using the original package contract. mixed = self.root/'mixed.zip' FakeStore.data = data export(target, mixed, self.settings, FakeStore) self.assertEqual(verify_archive(mixed)[0]['snapshot'], snapshot) self.restore_package(mixed, self.root/'fresh', {}) self.assertEqual(load(self.root/'fresh'), actual) self.assertEqual(self.restore_package(second, target, data, 'article'), (snapshot, 1)) # The independent verifier and online reader must follow each library's # pinned collection, rather than inventing a shared collection name. from argparse import Namespace from src.step2_knowledge_ingest import ingest from src.step3_question_answer.retrieval import Retriever, RetrievalError FakeStore.data = data with patch.object(ingest, 'config', return_value=self.settings), \ patch.object(ingest, 'model_fingerprint', return_value={}), \ patch.object(ingest, 'model_matches_manifest', return_value=True), \ patch.object(ingest, 'Store', FakeStore): ingest.execute(Namespace(prepared=target, env_file=self.root/'missing.env', verify_only=True, dry_run=False)) with patch('src.step3_question_answer.retrieval.config', return_value=self.settings): reader = Retriever(self.root/'missing.env', target/'ingestion/manifest.json') self.assertEqual(reader.manifest['databases']['article'], current['databases']['article']) bad = copy.deepcopy(current) bad['databases']['article']['snapshot'] = '../wrong' (target/'ingestion/manifest.json').write_text(canonical(bad), encoding='utf-8') with self.assertRaises(RetrievalError): Retriever(self.root/'missing.env', target/'ingestion/manifest.json') def test_scope_and_model_fail_before_writing(self): full = self.version(self.six_records(), 'full') single = self.version([make_record('general', 10)], 'single', database='article') wrong_model = self.version([make_record('general', 11)], 'model', database='article', model_files={'weights.bin':'f'*64}) target, data = self.root/'target', {} self.restore_package(full, target, data) before = copy.deepcopy(data) original = (target/'ingestion/manifest.json').read_bytes() for package, library in ((single,None),(single,'policy'),(full,'article'),(single,'unknown'),(wrong_model,'article')): with self.subTest(library=library,package=package.name), self.assertRaises(ValueError): self.restore_package(package, target, data, library) self.assertEqual(data,before) self.assertEqual((target/'ingestion/manifest.json').read_bytes(),original) with self.assertRaisesRegex(ValueError,'single_library_requires_existing_snapshot'): self.restore_package(single,self.root/'empty',{},'article') def test_every_library_can_be_updated_without_touching_other_databases(self): rows = self.six_records() full = self.version(rows, 'six') for ordinal, original in enumerate(rows, 20): library = route(original) replacement = copy.deepcopy(original) replacement['id'] = 'k_' + format(ordinal,'024x') package = self.version([replacement], 'only-'+library, database=library) with self.subTest(library=library): target, data = self.root/('target-'+library), {} self.restore_package(full,target,data) before=copy.deepcopy(data) self.restore_package(package,target,data,library) actual=load(target) self.assertEqual([r for r in actual if route(r)==library],[replacement]) self.assertEqual({r['id']:r for r in actual if route(r)!=library}, {r['id']:r for r in rows if route(r)!=library}) self.assertEqual({k:v for k,v in data.items() if k!='demo_'+library}, {k:v for k,v in before.items() if k!='demo_'+library}) def test_import_refuses_concurrent_publication(self): from src.step2_knowledge_ingest.ingest import publication_lock target = self.root/'target' with publication_lock(target): with self.assertRaisesRegex(ValueError,'another_knowledge_operation_is_active'): self.restore_package(self.root/'not-needed.zip',target,{}) self.assertFalse(target.exists()) def test_remote_failure_and_publication_failure_preserve_live_snapshot(self): first = self.version(self.six_records(), 'first') second = self.version([make_record('general', 10)], 'single', database='article') target, data = self.root/'target', {} self.restore_package(first,target,data) original = {p.relative_to(target):p.read_bytes() for p in target.rglob('*') if p.is_file()} class Failing(FakeStore): def verify(self,*args): raise ValueError('synthetic_remote_failure') with self.assertRaisesRegex(ValueError,'synthetic_remote_failure'): self.restore_package(second,target,data,'article',Failing) self.assertEqual(original,{p.relative_to(target):p.read_bytes() for p in target.rglob('*') if p.is_file()}) rename = Path.rename def fail_publish(path,dest): if path.name=='prepared' and '-import-' in str(path.parent): raise OSError('synthetic_publish_failure') return rename(path,dest) with patch.object(Path,'rename',fail_publish), self.assertRaisesRegex(OSError,'synthetic_publish_failure'): self.restore_package(second,target,data,'article') self.assertEqual(original,{p.relative_to(target):p.read_bytes() for p in target.rglob('*') if p.is_file()}) self.restore_package(second,target,data,'article') def test_cross_library_id_collision_and_corrupt_raw_are_rejected(self): rows=self.six_records() first=self.version(rows,'first') colliding=make_record('general',1) # ID already belongs to policy. single=self.version([colliding],'collision',database='article') target,data=self.root/'target',{} self.restore_package(first,target,data) before=copy.deepcopy(data) with self.assertRaisesRegex(ValueError,'invalid_id_or_kind'): self.restore_package(single,target,data,'article') self.assertEqual(data,before) # Rehashing the envelope cannot legitimize raw fields differing from # the enhanced record that will actually be indexed. from zipfile import ZipFile import hashlib with ZipFile(first) as z: members={n:z.read(n) for n in z.namelist()} name='prepared/policy.jsonl' raw=json.loads(members[name]);raw['title']='tampered' members[name]=(canonical(raw)+'\n').encode('utf-8') meta=json.loads(members['archive.json']) meta['sha256'][name]=hashlib.sha256(members[name]).hexdigest() members['archive.json']=canonical(meta).encode('utf-8') bad=self.root/'bad.zip' with ZipFile(bad,'w') as z: for n,b in members.items():z.writestr(n,b) with self.assertRaisesRegex(ValueError,'archive_original_mismatch'): verify_archive(bad) if __name__ == '__main__': unittest.main()