group_xinghuo_jinrong/tests/test_wave6_template_service.py

"""模板缓存（D-06 第二层）单元测试。"""
import unittest

from app.service.template_service import QueryTemplate, TemplateService


CUSTOMER_COUNT = QueryTemplate(
    template_key="customer_total_count",
    template_sql="SELECT COUNT(*) AS customer_count FROM core_customer",
    params_schema={"match_phrases": ["平台客户总数是多少"]},
    tags=["客户", "总数"],
)

SUBSCRIBE_DAYS = QueryTemplate(
    template_key="subscribe_amount_recent_days",
    template_sql=(
        "SELECT COALESCE(SUM(amount), 0) AS subscribe_total FROM core_trade "
        "WHERE trade_type = 'subscribe' "
        "AND trade_date >= DATE_SUB(CURDATE(), INTERVAL :days DAY)"
    ),
    params_schema={
        "match_all": ["申购", "金额"],
        "params": [{"name": "days", "placeholder": ":days", "extract": "recent_days", "default": 30}],
    },
    tags=["申购", "金额"],
)


class TestTemplateService(unittest.TestCase):
    def setUp(self):
        self.svc = TemplateService(templates=[CUSTOMER_COUNT, SUBSCRIBE_DAYS])

    def test_match_customer_count(self):
        hit = self.svc.try_render("平台客户总数是多少", "full", [], ["analyst"])
        self.assertIsNotNone(hit)
        sql, key = hit
        self.assertEqual(key, "customer_total_count")
        self.assertIn("COUNT(*)", sql)

    def test_match_subscribe_with_days(self):
        hit = self.svc.try_render("近7天申购金额总额", "full", [], ["analyst"])
        self.assertIsNotNone(hit)
        sql, key = hit
        self.assertEqual(key, "subscribe_amount_recent_days")
        self.assertIn("INTERVAL 7 DAY", sql)
        self.assertNotIn(":days", sql)

    def test_default_days_when_unspecified(self):
        hit = self.svc.try_render("申购金额汇总", "full", [], ["analyst"])
        self.assertIsNotNone(hit)
        sql, _ = hit
        self.assertIn("INTERVAL 30 DAY", sql)

    def test_no_match_returns_none(self):
        self.assertIsNone(self.svc.try_render("随便问问", "full", [], ["analyst"]))

    def test_match_nav_max_min_dual_rows_sql(self):
        tpl = QueryTemplate(
            template_key="product_nav_max_min_latest",
            template_sql=(
                "(SELECT 'max' AS extremum, n.product_id FROM core_product_nav n LIMIT 1) "
                "UNION ALL (SELECT 'min' AS extremum, n.product_id FROM core_product_nav n LIMIT 1)"
            ),
            params_schema={"match_phrases": ["目前净值最高和最低的两个产品"]},
            tags=["净值", "涨幅"],
        )
        svc = TemplateService(templates=[tpl])
        hit = svc.try_render("目前净值最高和最低的两个产品", "full", [], ["analyst"])
        self.assertIsNotNone(hit)
        sql, key = hit
        self.assertEqual(key, "product_nav_max_min_latest")
        self.assertIn("UNION ALL", sql.upper())

    def test_match_nav_chg_template_by_exact_phrase(self):
        nav_tpl = QueryTemplate(
            template_key="product_nav_max_min_latest",
            template_sql="SELECT 1 AS nav_tpl",
            params_schema={"match_phrases": ["目前净值最高和最低的两个产品"]},
        )
        chg_tpl = QueryTemplate(
            template_key="product_nav_chg_max_min_latest",
            template_sql="SELECT 1 AS chg_tpl",
            params_schema={"match_phrases": ["最新净值日涨跌幅最高和最低各一只产品"]},
        )
        svc = TemplateService(templates=[nav_tpl, chg_tpl])
        hit_nav = svc.try_render("目前净值最高和最低的两个产品", "full", [], ["analyst"])
        hit_chg = svc.try_render("最新净值日涨跌幅最高和最低各一只产品", "full", [], ["analyst"])
        self.assertEqual(hit_nav[1], "product_nav_max_min_latest")
        self.assertEqual(hit_chg[1], "product_nav_chg_max_min_latest")

    def test_phrase_template_does_not_fuzzy_match(self):
        tpl = QueryTemplate(
            template_key="product_nav_max_min_latest",
            template_sql="SELECT 1",
            params_schema={"match_phrases": ["目前净值最高和最低的两个产品"]},
        )
        svc = TemplateService(templates=[tpl])
        self.assertIsNone(svc.try_render("净值最高和最低的产品", "full", [], ["analyst"]))

    def test_list_published_prompts(self):
        tpl = QueryTemplate(
            template_key="x",
            template_sql="SELECT 1",
            params_schema={
                "match_phrases": ["平台客户总数是多少"],
                "prompt_label": "客户总数",
                "featured": True,
            },
        )
        svc = TemplateService(templates=[tpl])
        items = svc.list_published_prompts(["analyst"])
        self.assertEqual(len(items), 1)
        self.assertEqual(items[0]["question"], "平台客户总数是多少")

    def test_self_domain_injects_customer_id(self):
        tpl = QueryTemplate(
            template_key="self_trade_count",
            template_sql=(
                "SELECT COUNT(*) AS cnt FROM core_trade "
                "WHERE customer_id = ':customer_id'"
            ),
            params_schema={
                "match_all": ["交易", "多少"],
                "params": [
                    {"name": "customer_id", "placeholder": ":customer_id", "extract": "scope_customer", "default": ""}
                ],
            },
            tags=[],
        )
        svc = TemplateService(templates=[tpl])
        hit = svc.try_render("我有多少笔交易", "self", ["CUST-9527"], ["customer"])
        self.assertIsNotNone(hit)
        sql, _ = hit
        self.assertIn("CUST-9527", sql)


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