kg-agent/scripts/test_runtime_workflow.py

"""Offline end-to-end refresh contracts and generic schema query regressions."""
from __future__ import annotations
import copy
import hashlib
import json
import os
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
from openpyxl import Workbook

from step4_web import data_update as update
from step2_graph_building import runtime
from step3_qa_agent.retrieval.production_query import compile_query


def book(path, sheet, rows):
    path.parent.mkdir(parents=True,exist_ok=True)
    wb=Workbook()
    ws=wb.active
    ws.title=sheet
    for row in rows: ws.append(row)
    wb.save(path)
    wb.close()


class FakeDriver:
    def __init__(self):
        self.nodes={}
        self.edges={}
    def execute_query(self,q,**p):
        b=p.get('build')
        if 'SET n += row.props' in q:
            self.nodes.setdefault(b,set()).update(r['id'] for r in p['rows'])
        elif 'SET r._kg_method' in q:
            self.edges.setdefault(b,set()).update((r['source'],r['target'],p['rule']) for r in p['rows'])
        elif 'count(DISTINCT n)' in q:
            return SimpleNamespace(records=[{'nodes':len(self.nodes.get(b,set())), 'edges':len(self.edges.get(b,set()))}])
        elif 'DETACH DELETE n' in q:
            keys=list(self.nodes)
            for k in keys:
                if ('<>' in q and k!=b) or ('<>' not in q and k==b):
                    self.nodes.pop(k,None); self.edges.pop(k,None)
        return SimpleNamespace(records=[])


class UpdateTests(unittest.TestCase):
    def setUp(self):
        self.tmp=tempfile.TemporaryDirectory()
        self.root=Path(self.tmp.name)/'project'
        self.root.mkdir()
        self.templates=Path(self.tmp.name)/'shared templates/metadata'
        self.relation_dir=Path(self.tmp.name)/'shared relations'
        self.templates.mkdir(parents=True)
        self.write_template('人员信息',['工号','姓名'])
        self.write_template('项目信息',['项目编号','负责人编号'])
        self.write_relation()
        self.patches=[patch.object(update,'cleanup_answer_cache_for_release',return_value={'removed': 0}),
                      patch.object(update,'ROOT',self.root),
                      patch.object(update,'MANIFEST',self.root/'output/current_release.json'),
                      patch.object(runtime,'ROOT',self.root),
                      patch.object(runtime,'MANIFEST',self.root/'output/current_release.json'),
                      patch.dict(os.environ,{'STEP2_RELATION_SIMILARITY_THRESHOLD':'0.8',
                                             'METADATA_TEMPLATE_DIR':str(self.templates),
                                             'RELATION_DIR':str(self.relation_dir)})]
        for p in self.patches:p.start()
        self.driver=FakeDriver()
        self.fetch_calls=0
        self.assess_calls=0
        self.llm_calls=0

    def tearDown(self):
        for p in reversed(self.patches):p.stop()
        self.tmp.cleanup()

    def write_template(self,name,fields):
        book(self.templates/(name+'.xlsx'),'字段来源',
            [['模板字段','DMS模型','DMS字段','含义']]+[[f,'模型','c_'+str(i),'说明'] for i,f in enumerate(fields)])

    def write_relation(self,extra=False):
        rows=[['起点','终点','连接字段','判断方法','边名'],
              ['人员信息','项目信息','工号-负责人编号','相等','负责']]
        if extra:rows.append(['其他节点','项目信息','标识-项目编号','相等','关联'])
        book(self.relation_dir/'relation.xlsx','Sheet1',rows)

    def fetch(self,*,template_dir,output_dir,**kwargs):
        self.fetch_calls+=1
        for t in update.load_templates(template_dir):
            fields=[f.template_field for f in t.fields]
            book(output_dir/(t.name+'.xlsx'),'数据',[fields,['K']*len(fields)])

    def assess(self,*,template_dir,production_dir,analysis_dir,**kwargs):
        self.assess_calls+=1
        datasets=[]
        for t in update.load_templates(template_dir):
            source=production_dir/(t.name+'.xlsx')
            datasets.append({'dataset':t.name,'entity_identity':{'key_fields':[t.fields[0].template_field]},
                             'source_sha256':hashlib.sha256(source.read_bytes()).hexdigest()})
        analysis_dir.mkdir(parents=True)
        (analysis_dir/'qa_data_context.json').write_text(json.dumps({'raw_values_sent_to_llm':False,'datasets':datasets}),encoding='utf-8')

    def analyze(self,system,user):
        self.llm_calls+=1
        p=json.loads(user)
        return {k:[{'id':x['id'],'description':'测试说明'} for x in p[k]] for k in ('nodes','relations')}

    def run_update(self,mode='auto',**kwargs):
        return update.update_data(mode,fetcher=self.fetch,assessor=kwargs.get('assessor',self.assess),
                                  analyzer=self.analyze,driver=self.driver)

    def test_cache_cleanup_only_after_successful_publication(self):
        def cleaned(version):
            manifest = json.loads(update.MANIFEST.read_text(encoding='utf-8'))
            self.assertEqual(manifest['data_version'], version)
            return {'warning': 'Redis unavailable; retry pending'}
        with patch.object(update, 'cleanup_answer_cache_for_release', side_effect=cleaned) as cleanup:
            report = self.run_update()
            self.assertTrue(report['ok'])
            self.assertIn('warning', report['answer_cache_cleanup'])
            cleanup.assert_called_once_with(report['data_version'])
        with patch.object(update, 'cleanup_answer_cache_for_release') as cleanup:
            with patch.object(self, 'fetch', side_effect=ValueError('fetch failed')):
                with self.assertRaises(ValueError):
                    self.run_update()
            cleanup.assert_not_called()

    def test_dms_updates_assessment_but_preserves_schema_bytes(self):
        first=self.run_update()
        before=runtime.release_snapshot()
        schema_bytes=Path(before['schema_file']).read_bytes()
        display_bytes=Path(before['display_schema_file']).read_bytes()
        second=self.run_update('dms')
        after=runtime.release_snapshot()
        self.assertNotEqual(first['data_version'],second['data_version'])
        self.assertEqual(before['schema_version'],after['schema_version'])
        self.assertNotEqual(before['schema_file'],after['schema_file'])
        self.assertEqual(Path(after['schema_file']).parent.name, second['data_version'])
        self.assertEqual(Path(after['schema_file']).read_bytes(),schema_bytes)
        self.assertEqual(Path(after['display_schema_file']).read_bytes(),display_bytes)
        self.assertEqual(self.assess_calls,2)
        self.assertEqual(self.llm_calls,1)
        self.assertTrue(second['published'])
        self.assertFalse(second['schema_updated'])

    def test_retention_across_dms_updates_and_failures(self):
        import contextlib
        import io
        with contextlib.redirect_stdout(io.StringIO()):
            first=self.run_update('templates')
            first_schema=Path(runtime.release_snapshot()['schema_file']).read_bytes()
            for _ in range(11):
                latest=self.run_update('dms')
            snapshot=runtime.release_snapshot()
            base=self.root/'.runtime/releases'
            self.assertEqual(len(list(base.iterdir())),10)
            self.assertFalse((base/first['data_version']).exists())
            self.assertEqual(Path(snapshot['schema_file']).read_bytes(),first_schema)
            self.assertEqual(snapshot['schema_version'],first['schema_version'])
            self.assertRegex(latest['data_version'],r'^\d{8}_\d{6}_\d{6}_[0-9a-f]{32}$')
            self.assertTrue(latest['created_at'].endswith('+08:00'))
            self.assertTrue(latest['finished_at'].endswith('+08:00'))
            def broken(**kwargs): raise ValueError('synthetic failure')
            for _ in range(11):
                with self.assertRaises(ValueError): self.run_update('dms',assessor=broken)
            self.assertEqual(len(list(base.iterdir())),10)
            self.assertEqual(runtime.release_snapshot()['data_version'],latest['data_version'])
            for field in ('schema_file','display_schema_file','production_dir','assessment_file'):
                self.assertTrue(Path(runtime.release_snapshot()[field]).exists())
            reports=[json.loads((d/'update_report.json').read_text(encoding='utf-8')) for d in base.iterdir()]
            self.assertTrue(any(not r['ok'] and 'retention' in r for r in reports))

    def test_added_node_and_relation_require_no_code_or_key_file(self):
        self.run_update()
        old=runtime.release_snapshot()
        self.write_template('其他节点',['标识','其他名称'])
        self.write_relation(extra=True)
        result=self.run_update()
        latest=runtime.release_snapshot()
        schema=json.loads(Path(latest['schema_file']).read_text(encoding='utf-8'))
        self.assertEqual(len(schema['nodes']),3)
        self.assertEqual(len(schema['relations']),2)
        self.assertTrue(result['schema_updated'])
        self.assertEqual(self.assess_calls,2)
        self.assertNotEqual(old['schema_version'],latest['schema_version'])
        self.assertFalse((self.root/'data/graph_keys.json').exists())

    def test_wrong_mode_and_failure_leave_current_release(self):
        self.run_update()
        old=(self.root/'output/current_release.json').read_bytes()
        self.write_template('其他节点',['标识'])
        with self.assertRaises(ValueError):self.run_update('dms')
        self.assertEqual(self.fetch_calls,1)
        def broken(**kwargs):raise ValueError('synthetic failure')
        with self.assertRaises(ValueError):self.run_update('templates',assessor=broken)
        self.assertEqual((self.root/'output/current_release.json').read_bytes(),old)
        self.assertEqual(len(self.driver.nodes),1)

    def test_missing_keys_warn_and_publish_in_templates_and_dms_modes(self):
        fetch = self.fetch
        def with_missing(**kwargs):
            fetch(**kwargs)
            book(kwargs['output_dir'] / '人员信息.xlsx', '数据',
                 [['工号', '姓名'], ['K', '有效人员'], [None, '缺编号'], [' ', '空白编号']])
        self.fetch = with_missing
        for mode in ('templates', 'dms'):
            with self.subTest(mode=mode), self.assertLogs(
                    'step2_graph_building.graph.production', level='WARNING'):
                result = self.run_update(mode)
            self.assertTrue(result['ok'])
            self.assertTrue(result['published'])
            self.assertEqual(result['skipped_rows'], 2)
            self.assertEqual(result['graph'], {'nodes': 2, 'edges': 1})
            self.assertEqual([w['row'] for w in result['warnings']], [3, 4])
            report_path = self.root / '.runtime/releases' / result['data_version'] / 'update_report.json'
            self.assertEqual(json.loads(report_path.read_text(encoding='utf-8'))['warnings'], result['warnings'])
            snapshot = runtime.release_snapshot()
            self.assertEqual(snapshot['data_counts']['人员信息']['skipped_rows'], 2)
            self.assertEqual(snapshot['data_counts']['人员信息']['merged_rows'], 0)

    def test_explicit_template_refresh_always_refreshes_assessment(self):
        self.run_update()
        self.run_update('templates')
        self.assertEqual((self.fetch_calls,self.assess_calls,self.llm_calls),(2,2,2))

    def test_missing_llm_identity_rejected_without_db_change(self):
        with self.assertRaises(ValueError):
            update.assessment_keys({'datasets':[{'dataset':'人员信息','entity_identity':{'key_fields':[]}}]})


class QueryTests(unittest.TestCase):
    def setUp(self):
        self.schema={'nodes':[{'id':'任意类别','attributes':['编号','名称','数量'],'multivalue_fields':{}},
                              {'id':'另一类别','attributes':['编号','描述']}],
                     'relations':[{'id':'r1','source':'任意类别','target':'另一类别','type':'关联'}]}
        self.spec={'nodes':[{'alias':'n','type':'任意类别'}], 'filters':[],
                   'select':[], 'aggregates':[{'alias':'n','op':'count','as':'总数'}]}
    def test_count_and_literal_parameterization(self):
        self.spec['filters']=[{'alias':'n','field':'名称','op':'contains','value':"x') DETACH DELETE n //"}]
        result=compile_query(self.spec,self.schema,'version')
        self.assertIn('count(DISTINCT n._kg_id)',result.query)
        self.assertNotIn('DETACH',result.query)
        self.assertEqual(result.params['v0'],"x') DETACH DELETE n //")
        self.assertIn('valueType',result.query)
        self.assertEqual(result.params['build'],'version')

    def test_unknown_fields_relations_and_unbounded_input_rejected(self):
        for patch_value in [{'nodes':[{'alias':'n','type':'不存在'}]},
                            {'limit':100000}, {'cypher':'MATCH(n) DELETE n'}]:
            spec={**self.spec,**patch_value}
            with self.assertRaises(ValueError):compile_query(spec,self.schema,'b')
        spec={**self.spec,'filters':[{'alias':'n','field':'私有字段','op':'eq','value':'x'}]}
        with self.assertRaises(ValueError):compile_query(spec,self.schema,'b')

    def test_dynamic_relationship_direction_and_dedup(self):
        self.spec['nodes'].append({'alias':'m','type':'另一类别'})
        self.spec['relations']=[{'source':'n','target':'m','id':'r1'}]
        result=compile_query(self.spec,self.schema,'b')
        self.assertIn('r0:',result.query)
        self.assertEqual(result.params['rule0'],'r1')
        self.spec['relations'][0].update(source='m',target='n')
        with self.assertRaises(ValueError):compile_query(self.spec,self.schema,'b')

    def test_disconnected_join_and_multivalue_sum_rejected(self):
        self.spec['nodes'].append({'alias':'m','type':'另一类别'})
        with self.assertRaises(ValueError):compile_query(self.spec,self.schema,'b')
        self.spec['nodes'].pop()
        self.schema['nodes'][0]['multivalue_fields']={'数量':1}
        self.spec['aggregates']=[{'alias':'n','field':'数量','op':'sum','as':'合计'}]
        with self.assertRaises(ValueError):compile_query(self.spec,self.schema,'b')


if __name__=='__main__':unittest.main(verbosity=2)