kg-agent/scripts/test_step2_production.py

"""Offline Step2 regressions with synthetic workbooks and fake external services."""
from __future__ import annotations

import json
import os
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch

import numpy as np
from openpyxl import Workbook

from step2_graph_building.graph.production import (
    Rule, identifier, load_rules, match_rule, prepare_graph, read_production, write_graph,
)
from step2_graph_building.meta.production import build_schema, export_schema, load_production_schema
from step2_graph_building.production_pipeline import build_production_graph


def workbook(path, title, headers, rows):
    wb = Workbook()
    ws = wb.active
    ws.title = title
    ws.append(headers)
    for row in rows:
        ws.append(row)
    wb.save(path)
    wb.close()


def fake_analyzer(system, user):
    data = json.loads(user)
    return {kind: [{'id': item['id'], 'description': '合成数据的语义说明'} for item in data[kind]]
            for kind in ('nodes', 'relations')}


class FakeDriver:
    def __init__(self, fail=False):
        self.calls = []
        self.nodes, self.edges = set(), set()
        self.fail = fail

    def execute_query(self, query, **params):
        self.calls.append((query, params))
        if 'SET n += row.props' in query:
            self.nodes.update(r['id'] for r in params['rows'])
        if 'SET r._kg_method' in query:
            if self.fail:
                raise RuntimeError('synthetic database failure')
            self.edges.update((r['source'], r['target'], params['rule']) for r in params['rows'])
        if 'count(DISTINCT n)' in query:
            return SimpleNamespace(records=[{'nodes': len(self.nodes), 'edges': len(self.edges)}])
        return SimpleNamespace(records=[])


class ProductionGraphTests(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.root = Path(self.tmp.name)
        self.templates = self.root / 'templates'
        self.production = self.root / 'production'
        self.output = self.root / 'output'
        self.templates.mkdir()
        self.production.mkdir()
        self.fields = {'人员信息': ['工号', '姓名', '服务项目'], '项目信息': ['项目编号', '项目名称', '上级编号']}
        for name, fields in self.fields.items():
            workbook(self.templates / f'{name}.xlsx', '字段来源',
                     ['模板字段', 'DMS模型', 'DMS字段', '含义'],
                     [(f, '模型', f'c_{i}', '') for i, f in enumerate(fields)])
        self.people = [('A', '合成人员甲', '北园 / 南园'), ('B', '合成人员乙', None),
                       ('C', '合成人员丙', '园'), ('D', '合成人员丁', '')]
        self.projects = [('P1', '北园', None), ('P2', '南园', 'P1'), ('P3', '东园', '')]
        self.write_data()
        self.relations = self.root / 'relation.xlsx'
        workbook(self.relations, '关系', ['起点', '终点', '连接字段', '判断方法', '边名'],
                 [('人员信息', '项目信息', '服务项目-项目名称', '包含', '服务于'),
                  ('项目信息', '项目信息', '上级编号-项目编号', '相等', '父项目为')])
        self.keys = {'人员信息': ['工号'], '项目信息': ['项目编号']}
        self.key_path = self.root / 'keys.json'
        self.key_path.write_text(json.dumps(self.keys), encoding='utf-8')
        self.env = patch.dict(os.environ, {'STEP2_RELATION_SIMILARITY_THRESHOLD': '0.8'})
        self.env.start()

    def tearDown(self):
        self.env.stop()
        self.tmp.cleanup()

    def write_data(self):
        for name, rows in [('人员信息', self.people), ('项目信息', self.projects)]:
            workbook(self.production / f'{name}.xlsx', '数据', self.fields[name], rows)

    def graph(self):
        return prepare_graph(self.templates, self.production, self.relations, keys=self.keys, threshold=.8)

    def pipeline(self, **kwargs):
        return build_production_graph(template_dir=self.templates, production_dir=self.production,
            relation_path=self.relations, keys_path=self.key_path, output_dir=self.output, **kwargs)

    def test_exact_contains_direction_and_blank(self):
        graph = self.graph()
        self.assertEqual([s['count'] for s in graph['relation_stats']], [5, 1])
        rule = graph['rules'][1]
        project_ids = {n['properties']['项目编号']: n['id'] for n in graph['nodes']['项目信息']}
        self.assertEqual(graph['edges'][rule.id], [{'source': project_ids['P2'], 'target': project_ids['P1'], 'score': None}])
        self.assertEqual(graph['relation_stats'][0]['unmatched_source_nodes'], 2)
        rule = Rule('A', 'B', 'v', 'v', '相等', 'r', 'test')
        nodes = {'A': [{'id': 'a', 'properties': {'v': ' x '}}],
                 'B': [{'id': 'b', 'properties': {'v': 'x'}}]}
        self.assertEqual(list(match_rule(rule, nodes, .8, None)), [])

    def test_merge_preserves_provenance_and_identity_across_reorder(self):
        original = self.graph()
        self.people.append(self.people[0])
        self.write_data()
        merged = self.graph()
        self.assertEqual(merged['node_stats']['人员信息']['source_rows'], 5)
        self.assertEqual(merged['node_stats']['人员信息']['node_count'], 4)
        self.assertEqual(merged['nodes']['人员信息'][0]['rows'], [2, 6])
        self.people.reverse()
        self.write_data()
        reordered = self.graph()
        self.assertEqual({n['id'] for n in original['nodes']['人员信息']},
                         {n['id'] for n in reordered['nodes']['人员信息']})

    def test_different_values_merge_without_losing_relationships(self):
        self.people += [('A', '不同属性', '北园'), ('C', '冲突', '园'), ('B', '合成人员乙', '东园')]
        self.write_data()
        graph = self.graph()
        stats = graph['node_stats']['人员信息']
        self.assertEqual(stats['node_count'], 4)
        self.assertEqual(stats['merged_rows'], 3)
        self.assertEqual(stats['merged_key_groups'], 3)
        person = next(n for n in graph['nodes']['人员信息'] if n['properties']['工号'] == 'A')
        self.assertEqual(set(person['properties']['服务项目']), {'北园 / 南园', '北园'})
        self.assertEqual(person['rows'], [2, 6])
        self.assertEqual(graph['relation_stats'][0]['count'], 6)
        driver = FakeDriver()
        result = self.pipeline(driver=driver, analyzer=fake_analyzer)
        self.assertTrue(result['ok'])
        schema = load_production_schema(self.output / 'meta_graph_schema.json')
        self.assertEqual(schema['nodes'][0]['multivalue_fields']['服务项目'], 1)
        self.people.reverse()
        self.write_data()
        again = self.graph()
        original = {n['id']: n['properties'] for n in graph['nodes']['人员信息']}
        self.assertEqual(original, {n['id']: n['properties'] for n in again['nodes']['人员信息']})

    def test_multivalued_exact_and_score_deduplicate_edges(self):
        from step2_graph_building.graph.production import neo4j_properties, merge_value
        self.assertEqual(merge_value(None, 'x'), 'x')
        self.assertEqual(merge_value('x', ''), 'x')
        self.assertEqual(merge_value('x', 'x'), 'x')
        self.assertEqual(merge_value(0, None), 0)
        mixed = neo4j_properties({'x': [1, '1']})
        self.assertEqual(json.loads(mixed['_kg_mixed_values_json']), {'x': [1, '1']})
        rule = Rule('A', 'B', 'v', 'v', '相等', 'r', 'test')
        nodes = {'A': [{'id': 'a', 'properties': {'v': ['x', 'y']}}],
                 'B': [{'id': 'b', 'properties': {'v': ['y', 'z']}}]}
        self.assertEqual(len(list(match_rule(rule, nodes, .8, None))), 1)
        self.people += [('A', '合成人员甲', '北园'), ('A', '合成人员甲', '南园')]
        self.write_data()
        workbook(self.relations, '关系', ['起点', '终点', '连接字段', '判断方法', '边名'],
                 [('人员信息', '项目信息', '服务项目-项目名称', '评分', '相似于')])
        vectors = {'北园 / 南园': [.9, .1], '北园': [1, 0], '南园': [.85, .2],
                   '园': [0, 1], '东园': [0, 1]}
        encoder = SimpleNamespace(encode=lambda texts: np.array([vectors[t] for t in texts]))
        graph = prepare_graph(self.templates, self.production, self.relations,
                              keys=self.keys, threshold=.8, encoder=encoder)
        person = next(n['id'] for n in graph['nodes']['人员信息'] if n['properties']['工号'] == 'A')
        edges = [e for e in graph['edges'][graph['rules'][0].id] if e['source'] == person]
        self.assertEqual(len(edges), 2)
        self.assertTrue(all(abs(e['score'] - 1) < 1e-12 for e in edges))

    def test_missing_keys_warn_skip_and_preserve_valid_graph(self):
        original = self.graph()
        self.people += [(None, '缺编号', '园'), ('   ', '空白编号', '园'),
                        self.people[0], (None, None, None)]
        self.write_data()
        before = (self.production / '人员信息.xlsx').read_bytes()
        with self.assertLogs('step2_graph_building.graph.production', level='WARNING') as logs:
            graph = self.graph()
        self.assertEqual(graph['nodes'], original['nodes'] | {
            '人员信息': [dict(n, rows=[2, 8]) if n['properties']['工号'] == 'A' else n
                         for n in original['nodes']['人员信息']]})
        self.assertEqual(graph['edges'], original['edges'])
        self.assertEqual(graph['skipped_rows'], 2)
        self.assertEqual([w['row'] for w in graph['warnings']], [6, 7])
        self.assertTrue(all(w['fields'] == ['工号'] and w['action'] == 'skipped'
                            for w in graph['warnings']))
        stats = graph['node_stats']['人员信息']
        self.assertEqual((stats['source_rows'], stats['skipped_rows'], stats['merged_rows'], stats['node_count']),
                         (7, 2, 1, 4))
        self.assertIn('第 6 行', '\n'.join(logs.output))
        self.assertEqual((self.production / '人员信息.xlsx').read_bytes(), before)
        result = self.pipeline(check_only=True)
        self.assertTrue(result['ok'])
        saved = json.loads((self.output / 'step2_build_report.json').read_text(encoding='utf-8'))
        self.assertEqual(saved['warnings'], graph['warnings'])
        captured = []
        def analyzer(system, user):
            captured.append(user)
            return fake_analyzer(system, user)
        build_schema(graph, analyzer=analyzer, model='fake', build_id='test')
        self.assertNotIn(str(self.production), captured[0])
        self.assertNotIn('缺编号', captured[0])

    def test_composite_keys_only_report_missing_parts_and_allow_zero(self):
        self.keys['人员信息'] = ['工号', '姓名']
        self.people = [('A', None, '园'), (0, '零编号', '园')]
        self.write_data()
        graph = self.graph()
        self.assertEqual(graph['warnings'][0]['fields'], ['姓名'])
        self.assertEqual(graph['nodes']['人员信息'][0]['properties']['工号'], 0)
        self.assertEqual(graph['node_stats']['人员信息']['merged_rows'], 0)
        with self.assertRaises(ValueError):
            read_production(self.templates, self.production, keys={})
        with self.assertRaises(ValueError):
            read_production(self.templates, self.production, keys={**self.keys, '未知模板': ['编号']})

    def test_all_missing_keys_keeps_empty_template_without_false_merges(self):
        self.people = [(None, '缺编号', '园')]
        self.write_data()
        graph = self.graph()
        self.assertEqual(graph['nodes']['人员信息'], [])
        self.assertEqual(graph['node_stats']['人员信息']['merged_rows'], 0)
        self.assertEqual(graph['node_stats']['人员信息']['skipped_rows'], 1)
        self.assertEqual(graph['relation_stats'][0]['count'], 0)

    def test_bad_relations_and_schema_rejected(self):
        for row in [
            ('人员信息', '未知', '服务项目-项目名称', '包含', '服务于'),
            ('人员信息', '项目信息', '不存在-项目名称', '包含', '服务于'),
            ('人员信息', '项目信息', '服务项目-项目名称', '相似', '服务于'),
        ]:
            workbook(self.relations, '关系', ['起点', '终点', '连接字段', '判断方法', '边名'], [row])
            with self.assertRaises(ValueError):
                self.graph()
        workbook(self.production / '人员信息.xlsx', '数据', ['错误表头'], [('x',)])
        with self.assertRaises(ValueError):
            read_production(self.templates, self.production, keys=self.keys)

    def test_score_cosine_strict_threshold_all_matches(self):
        rule = Rule('A', 'B', 'v', 'v', '评分', '相似于', 'test')
        nodes = {'A': [{'id': 'a', 'properties': {'v': 'source'}}],
                 'B': [{'id': key, 'properties': {'v': key}} for key in ['equal', 'higher', 'lower', 'same']]}
        vectors = {'source': [1, 0], 'equal': [.8, .6], 'higher': [.9, .1],
                   'lower': [.7, .7], 'same': [1, 0]}
        encoder = SimpleNamespace(encode=lambda texts: np.array([vectors[t] for t in texts]))
        matches = list(match_rule(rule, nodes, .8, encoder))
        self.assertEqual({m['target'] for m in matches}, {'higher', 'same'})
        self.assertTrue(all(m['score'] > .8 for m in matches))
        with self.assertRaises(ValueError):
            prepare_graph(self.templates, self.production, self.relations, keys=self.keys, threshold=float('nan'))

    def test_meta_deterministic_topology_and_no_raw_values(self):
        graph = self.graph()
        captured = []
        def analyzer(system, user):
            captured.append(user)
            return fake_analyzer(system, user)
        schema = build_schema(graph, analyzer=analyzer, model='fake', build_id='test')
        self.assertNotIn('合成人员甲', captured[0])
        self.assertNotIn('北园 / 南园', captured[0])
        self.assertNotIn(str(self.root), captured[0])
        self.assertEqual(schema['nodes'][0]['attributes'], self.fields['人员信息'])
        self.assertEqual(schema['relations'][0]['count'], 5)
        export_schema(schema, self.output)
        self.assertEqual(load_production_schema(self.output / 'meta_graph_schema.json'), schema)
        display = json.loads((self.output / 'meta_graph_schema_display.json').read_text(encoding='utf-8'))
        self.assertEqual(len(display['nodes']), 2)
        self.assertEqual(len(display['relations']), 2)
        for response in ({}, {'nodes': [], 'relations': []}):
            with self.assertRaises(ValueError):
                build_schema(graph, analyzer=lambda *_: response, model='bad', build_id='test')
        def fabricated(system, user):
            result = fake_analyzer(system, user)
            result['nodes'][0]['attributes'] = ['invented']
            return result
        with self.assertRaises(ValueError):
            build_schema(graph, analyzer=fabricated, model='bad', build_id='test')

    def test_batched_database_order_idempotence_and_identifier_escape(self):
        graph = self.graph()
        driver = FakeDriver()
        self.assertEqual(write_graph(graph, driver, build_id='test', batch_size=2), {'nodes': 7, 'edges': 6})
        self.assertEqual(write_graph(graph, driver, build_id='test', batch_size=2), {'nodes': 7, 'edges': 6})
        first_edge = next(i for i,(q,_) in enumerate(driver.calls) if 'SET r._kg_method' in q)
        self.assertEqual(sum(len(p['rows']) for q,p in driver.calls[:first_edge] if 'SET n += row.props' in q), 7)
        self.assertEqual(identifier('a`b'), '`a``b`')

    def test_pipeline_publish_and_rollback(self):
        driver = FakeDriver()
        result = self.pipeline(driver=driver, analyzer=fake_analyzer, model='fake')
        self.assertTrue(result['ok'])
        self.assertTrue(result['schema_published'])
        failed = FakeDriver(fail=True)
        before = (self.output / 'meta_graph_schema.json').read_bytes()
        result = self.pipeline(driver=failed, analyzer=fake_analyzer, model='fake')
        self.assertFalse(result['ok'])
        self.assertTrue(result['staging_rolled_back'])
        self.assertEqual((self.output / 'meta_graph_schema.json').read_bytes(), before)
        self.assertFalse(any('WHERE n._kg_build <>' in q for q,_ in failed.calls))

    def test_ambiguous_activation_keeps_new_build_for_recovery(self):
        class UncertainDriver(FakeDriver):
            def execute_query(self, query, **params):
                result = super().execute_query(query, **params)
                if 'WHERE n._kg_build <>' in query:
                    raise RuntimeError('activation response lost')
                return result
        driver = UncertainDriver()
        result = self.pipeline(driver=driver, analyzer=fake_analyzer, model='fake')
        self.assertFalse(result['ok'])
        self.assertEqual(result['stage'], 'neo4j_activation')
        self.assertNotIn('staging_rolled_back', result)
        self.assertFalse(any('{_kg_build:$build}) DETACH DELETE' in q for q,_ in driver.calls))
        self.assertTrue((self.output / 'step2_schema_candidate.json').exists())

    def test_qa_schema_uses_production_names(self):
        from step2_graph_building.meta.schema import llm_schema_entities, llm_schema_relations
        schema = build_schema(self.graph(), analyzer=fake_analyzer, model='fake', build_id='test')
        with patch('step2_graph_building.meta.production.load_production_schema', return_value=schema):
            self.assertIn('人员信息', llm_schema_entities())
            self.assertIn('(人员信息)-[:服务于]->(项目信息)', llm_schema_relations())
            self.assertNotIn('有考勤', llm_schema_relations())

    def test_check_only_and_llm_failure_never_write_database(self):
        driver = FakeDriver()
        result = self.pipeline(check_only=True, driver=driver)
        self.assertTrue(result['ok'])
        self.assertFalse(driver.calls)
        result = self.pipeline(driver=driver, analyzer=lambda *_: {})
        self.assertFalse(result['ok'])
        self.assertEqual(result['stage'], 'meta_generation')
        self.assertFalse(driver.calls)


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