kg-agent/scripts/test_data_quality.py

"""Data quality API tests with two source files and overlapping row numbers."""
import json
import tempfile
import unittest
from datetime import datetime
from pathlib import Path
from unittest.mock import patch

from fastapi import FastAPI
from fastapi.testclient import TestClient
from openpyxl import Workbook

from step4_web import data_quality as quality


class DataQualityTests(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.root = Path(self.tmp.name)
        self.production = self.root / 'production'
        self.production.mkdir()
        self.snapshot = {'data_version': 'v1', 'production_dir': str(self.production)}
        self.report = {'data_version': 'v1', 'warnings': []}
        for name in ('人员证书', '项目信息'):
            wb = Workbook(); ws = wb.active; ws.title = '数据'
            ws.append(['编号', '名称', '日期', '数量', '备注'])
            ws.append([None, name + '<script>alert(1)</script>', datetime(2026, 9, 4), 0, None])
            ws.append(['valid', '正常数据不应暴露', None, 1, None])
            wb.save(self.production / (name + '.xlsx')); wb.close()
            self.report['warnings'].append({'type': 'missing_key', 'template': name,
                'file': 'ignored-old-path', 'row': 2, 'fields': ['编号'], 'sheet': '数据'})
        self.save()
        self.mock = patch.object(quality, 'release_snapshot', return_value=self.snapshot)
        self.mock.start()
        app = FastAPI(); app.include_router(quality.router)
        self.client = TestClient(app)

    def tearDown(self):
        self.client.close(); self.mock.stop(); self.tmp.cleanup()

    def save(self):
        (self.root / 'update_report.json').write_text(json.dumps(self.report), encoding='utf-8')

    def test_summary_counts_files_without_loading_raw_records(self):
        with patch.object(quality, 'load_workbook', side_effect=AssertionError('summary opened Excel')):
            response = self.client.get('/api/data-quality')
        self.assertEqual(response.status_code, 200)
        self.assertEqual(response.headers['cache-control'], 'no-store')
        self.assertEqual(response.json(), {'data_version':'v1', 'total':2,
                         'datasets': {'人员证书':1, '项目信息':1}})

    def test_details_separate_same_row_number_and_include_every_field(self):
        response = self.client.get('/api/data-quality?details=true&data_version=v1')
        self.assertEqual(response.status_code, 200)
        records = response.json()['records']
        self.assertEqual([r['file'] for r in records], ['人员证书.xlsx', '项目信息.xlsx'])
        for record in records:
            self.assertEqual(record['row'], 2)
            self.assertEqual(len(record['data']), 5)
            self.assertIsNone(record['data']['编号'])
            self.assertIsNone(record['data']['备注'])
            self.assertEqual(record['data']['数量'], 0)
            self.assertEqual(record['data']['日期'], '2026-09-04T00:00:00')
            self.assertTrue(record['data']['名称'].startswith(record['template']))
        self.assertNotIn('正常数据', response.text)
        self.assertNotIn(str(self.production), response.text)

    def test_version_change_before_and_during_read_requires_refresh(self):
        self.assertEqual(self.client.get('/api/data-quality?details=true&data_version=old').status_code, 409)
        with patch.object(quality, 'release_snapshot', side_effect=[self.snapshot, {**self.snapshot, 'data_version':'v2'}]):
            self.assertEqual(self.client.get('/api/data-quality?details=true').status_code, 409)

    def test_missing_source_keeps_problem_and_other_file_data(self):
        (self.production / '人员证书.xlsx').unlink()
        records = self.client.get('/api/data-quality?details=true').json()['records']
        self.assertIn('data_error', records[0])
        self.assertIsNotNone(records[1]['data'])

    def test_no_issues_and_missing_report_are_different(self):
        self.report['warnings'] = []; self.save()
        self.assertEqual(self.client.get('/api/data-quality?details=true').json()['records'], [])
        (self.root / 'update_report.json').unlink()
        self.assertEqual(self.client.get('/api/data-quality').status_code, 503)

    def test_report_version_mismatch_and_outside_source_rejected(self):
        self.report['data_version'] = 'old'; self.save()
        self.assertEqual(self.client.get('/api/data-quality').status_code, 503)
        self.report['data_version'] = 'v1'
        self.report['warnings'][0]['template'] = '../outside'; self.save()
        self.assertEqual(self.client.get('/api/data-quality?details=true').status_code, 503)


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