YG-Rules/tests/test_schema_descriptions.py

import io
import os
import tempfile
import unittest

from app.utils.llm import strip_thinking
from app.utils.schema_storage import SchemaStorage


MODULE = {
    "module_name": "1、银行账户",
    "fields": [
        {"name": "所属集团", "marker": "groupName", "required": "必填", "rule": "填写集团名称"},
        {"name": "开户单位", "marker": "accountOrg", "required": "必填", "rule": "填写账户所属单位"},
        {"name": "银行账号", "marker": "accountNo", "required": "必填", "rule": "填写银行账号"},
        {"name": "开户银行", "marker": "bankName", "required": "必填", "rule": "填写开户银行"},
        {"name": "币种", "marker": "currency", "required": "必填", "rule": "填写账户币种"},
        {"name": "账户余额", "marker": "balance", "required": "必填", "rule": "填写账户余额"},
    ],
}


class LLMThinkingCleanupTest(unittest.TestCase):
    def test_strip_complete_think_block(self):
        text = "<think>这里是推理过程</think>银行账户表用于管理账户基础信息。"
        self.assertEqual(strip_thinking(text), "银行账户表用于管理账户基础信息。")

    def test_strip_unclosed_think_block(self):
        self.assertEqual(strip_thinking("<think>只有推理过程，没有最终答案"), "")


class SchemaDescriptionSanitizingTest(unittest.TestCase):
    def setUp(self):
        self.storage = object.__new__(SchemaStorage)

    def test_extracts_final_quoted_description(self):
        response = (
            "好的，我先分析字段。\n"
            "\"银行账户表用于记录集团及成员单位银行账户基础信息，包含所属集团、开户单位、"
            "银行账号、开户银行、币种、账户余额等字段，支撑账户查询、余额校验和资金管理选表。\""
        )
        description = self.storage._sanitize_description(response, MODULE)

        self.assertNotIn("好的", description)
        self.assertIn("银行账户", description)
        self.assertIn("账户余额", description)

    def test_falls_back_when_only_thinking_is_returned(self):
        description = self.storage._sanitize_description("<think>仅有思考过程</think>", MODULE)

        self.assertNotIn("<think>", description)
        self.assertIn("银行账户", description)
        self.assertIn("所属集团", description)

    def test_cleans_historical_schema_response(self):
        data = {
            "processing_status": "done",
            "modules": [
                {**MODULE, "description": "<think>历史脏数据</think>"},
                {**MODULE, "module_name": "2、资金结算", "description": ""},
            ],
        }

        cleaned = self.storage._clean_descriptions_for_response(data)

        self.assertNotIn("<think>", cleaned["modules"][0]["description"])
        self.assertIn("银行账户", cleaned["modules"][0]["description"])
        self.assertIn("资金结算", cleaned["modules"][1]["description"])

    def test_parse_excel_preserves_physical_table_name(self):
        from openpyxl import Workbook

        workbook = Workbook()
        sheet = workbook.active
        sheet.append(["序号", "表名", "数据名称", "数据标记", "数据类型", "数据长度", "数据填写规则", "数据填写要求", "强弱校验"])
        sheet.append(["银行贷款"])
        sheet.append([1, "bank_loan", "贷款余额", "STANDARDCURRENCYBALANCE", "数值", "18,2", "填写贷款余额", "必填", "强校验"])

        stream = io.BytesIO()
        workbook.save(stream)
        stream.seek(0)

        modules = self.storage._parse_excel(stream)

        self.assertEqual(modules[0]["table_name"], "bank_loan")
        self.assertEqual(modules[0]["fields"][0]["marker"], "STANDARDCURRENCYBALANCE")

    def test_delete_file_removes_schema_json(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            schema_path = os.path.join(temp_dir, "schema.json")
            storage = SchemaStorage(schema_path)

            self.assertTrue(os.path.exists(schema_path))
            self.assertTrue(storage.delete_file())
            self.assertFalse(os.path.exists(schema_path))
            self.assertFalse(storage.delete_file())


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