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