kg-agent/scripts/test_data_analysis.py

"""Unit tests for production workbook profiling and QA context generation."""

from __future__ import annotations

import json
import threading
import tempfile
import unittest
from pathlib import Path

from openpyxl import Workbook

from step1_data_aggregation.data_analysis import analyze_production_data, find_candidate_keys
from step1_data_aggregation.template_mapping import DMS_FIELD, DMS_MODEL, MAPPING_SHEET, TEMPLATE_FIELD
from step3_qa_agent.agent.data_context import production_data_context_text


def write_template(path: Path) -> None:
    workbook = Workbook()
    sheet = workbook.active
    sheet.title = MAPPING_SHEET
    sheet.append([TEMPLATE_FIELD, DMS_MODEL, DMS_FIELD, "含义"])
    sheet.append(["姓名", "Group-People", "c_name", "员工姓名"])
    sheet.append(["部门", "Group-People", "c_org", "所属部门"])
    workbook.save(path)
    workbook.close()


def write_production(path: Path) -> None:
    workbook = Workbook()
    sheet = workbook.active
    sheet.title = "数据"
    sheet.append(["姓名", "部门"])
    sheet.append(["Alice", "工程部"])
    sheet.append(["Bob", None])
    workbook.save(path)
    workbook.close()


class DataAnalysisTests(unittest.TestCase):
    def test_profiles_without_sending_raw_values_and_builds_qa_context(self) -> None:
        with tempfile.TemporaryDirectory() as raw_dir:
            root = Path(raw_dir)
            template_dir = root / "templates"
            production_dir = root / "production"
            template_dir.mkdir()
            production_dir.mkdir()
            write_template(template_dir / "人员.xlsx")
            write_production(production_dir / "人员.xlsx")
            captured_user = ""

            def fake_analyzer(system: str, user: str):
                nonlocal captured_user
                captured_user = user
                self.assertIn("不得猜测", system)
                return {
                    "dataset_summary": "人员基础数据，部门字段存在缺失。",
                    "routing_terms": ["人员", "员工"],
                    "quality_assessment": {
                        "score": 75,
                        "level": "fair",
                        "strengths": ["姓名完整"],
                        "issues": [
                            {
                                "severity": "medium",
                                "field": "部门",
                                "issue": "存在缺失",
                                "evidence": "缺失率50%",
                                "recommendation": "补齐部门",
                            }
                        ],
                    },
                    "fields": [
                        {
                            "field": name,
                            "business_meaning": meaning,
                            "data_characteristics": "文本字段",
                            "quality_findings": [],
                            "qa_usage": {
                                "query_intents": ["筛选"],
                                "filterable": True,
                                "aggregatable": False,
                                "join_candidate": name == "姓名",
                                "cautions": [],
                            },
                        }
                        for name, meaning in (("姓名", "员工姓名"), ("部门", "所属部门"))
                    ],
                    "record_granularity": {
                        "one_row_represents": "一名员工",
                        "candidate_business_key": ["姓名"],
                        "cardinality_notes": ["员工与部门为 N:1"],
                        "confidence": "medium",
                        "evidence": ["姓名在观察数据中唯一，但仍需业务确认"],
                    },
                    "dataset_qa_guidance": {
                        "suitable_questions": ["按部门筛选人员"],
                        "unsupported_or_risky_questions": [],
                        "interpretation_notes": ["部门缺失会导致漏计"],
                    },
                }

            results = analyze_production_data(
                template_dir=template_dir,
                production_dir=production_dir,
                analyzer=fake_analyzer,
                model="fake-model",
            )
            self.assertEqual(len(results), 1)
            self.assertNotIn("Alice", captured_user)
            self.assertNotIn("Bob", captured_user)
            self.assertNotIn("c_name", captured_user)
            self.assertNotIn(str(production_dir), captured_user)
            payload = json.loads(Path(results[0].profile_file).read_text(encoding="utf-8"))
            self.assertFalse(payload["profile"]["privacy"]["raw_values_sent_to_llm"])
            self.assertEqual(payload["profile"]["fields"][1]["missing_rate"], 0.5)
            self.assertEqual(
                payload["profile"]["record_grain_profile"]["candidate_keys"][0]["fields"],
                ["姓名"],
            )
            context_path = production_dir / "analysis" / "qa_data_context.json"
            context = production_data_context_text(context_path, question="按部门筛选人员")
            self.assertIn("人员", context)
            self.assertIn("记录唯一性依据（候选业务主键）：姓名", context)
            self.assertIn("部门缺失会导致漏计", context)
            self.assertNotIn("Alice", context)

            self.assertEqual(
                production_data_context_text(context_path, question="项目到期"), ""
            )

    def test_finds_minimal_composite_candidate_key(self) -> None:
        candidates = find_candidate_keys(
            ["工号", "证书"],
            [("E001", "A"), ("E001", "B"), ("E002", "A")],
        )
        self.assertEqual(candidates[0]["fields"], ["工号", "证书"])


    def test_analyzes_independent_workbooks_concurrently(self) -> None:
        with tempfile.TemporaryDirectory() as raw_dir:
            root = Path(raw_dir)
            template_dir = root / "templates"
            production_dir = root / "production"
            template_dir.mkdir()
            production_dir.mkdir()
            for name in ("人员", "项目"):
                write_template(template_dir / f"{name}.xlsx")
                write_production(production_dir / f"{name}.xlsx")
            barrier = threading.Barrier(2, timeout=3)
            thread_names: set[str] = set()
            lock = threading.Lock()

            def fake_analyzer(system: str, user: str):
                del system, user
                with lock:
                    thread_names.add(threading.current_thread().name)
                barrier.wait()
                return {
                    "dataset_summary": "测试数据。",
                    "routing_terms": ["测试"],
                    "quality_assessment": {
                        "score": 80,
                        "level": "good",
                        "strengths": [],
                        "issues": [],
                    },
                    "fields": [
                        {
                            "field": name,
                            "business_meaning": meaning,
                            "data_characteristics": "文本",
                            "quality_findings": [],
                            "qa_usage": {
                                "query_intents": ["筛选"],
                                "filterable": True,
                                "aggregatable": False,
                                "join_candidate": name == "姓名",
                                "cautions": [],
                            },
                        }
                        for name, meaning in (("姓名", "员工姓名"), ("部门", "所属部门"))
                    ],
                    "record_granularity": {
                        "one_row_represents": "一条测试记录",
                        "candidate_business_key": ["姓名+部门"],
                        "cardinality_notes": [],
                        "confidence": "low",
                        "evidence": [],
                    },
                    "dataset_qa_guidance": {
                        "suitable_questions": [],
                        "unsupported_or_risky_questions": [],
                        "interpretation_notes": [],
                    },
                }

            results = analyze_production_data(
                template_dir=template_dir,
                production_dir=production_dir,
                analyzer=fake_analyzer,
                model="fake-model",
                max_workers=2,
            )
            self.assertEqual(len(results), 2)
            self.assertEqual(len(thread_names), 2)
            payload = json.loads(Path(results[0].profile_file).read_text(encoding="utf-8"))
            self.assertEqual(
                payload["llm_analysis"]["record_granularity"]["candidate_business_key"],
                ["姓名", "部门"],
            )


if __name__ == "__main__":
    unittest.main()