business-assistant/src/step2_knowledge_ingest/tests/test_portable.py

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