group_xinghuo_jinrong/tests/test_wave6_sql_guard.py

"""sql_guard 单元测试（Wave 6）。"""
import unittest

from app.service.sql_guard import SqlGuardError, extract_tables, inject_ownership, normalize_union_all_sql, validate


class TestSqlGuard(unittest.TestCase):
    def test_select_allowed(self):
        r = validate("SELECT risk_code, COUNT(*) AS c FROM core_customer GROUP BY risk_code", "full")
        self.assertTrue(r.allowed)

    def test_paren_union_select_allowed(self):
        sql = (
            "(SELECT 'max' AS extremum, product_id FROM core_product_nav ORDER BY daily_chg_pct DESC LIMIT 1) "
            "UNION ALL "
            "(SELECT 'min' AS extremum, product_id FROM core_product_nav ORDER BY daily_chg_pct ASC LIMIT 1)"
        )
        r = validate(sql, "full")
        self.assertTrue(r.allowed)

    def test_normalize_union_wraps_branches(self):
        raw = (
            "SELECT product_id FROM core_product_nav ORDER BY nav DESC LIMIT 1 UNION ALL "
            "SELECT product_id FROM core_product_nav ORDER BY nav ASC LIMIT 1"
        )
        fixed = normalize_union_all_sql(raw)
        self.assertTrue(fixed.startswith("("))
        self.assertIn("UNION ALL", fixed.upper())
        r = validate(fixed, "full")
        self.assertTrue(r.allowed)

    def test_customer_self_in_scope(self):
        r = validate(
            "SELECT COUNT(*) FROM core_holding WHERE customer_id='CUST-9527'",
            "self",
            ["CUST-9527"],
        )
        self.assertTrue(r.allowed)

    def test_customer_self_out_of_scope(self):
        with self.assertRaises(SqlGuardError) as cm:
            validate(
                "SELECT * FROM core_holding WHERE customer_id='CUST-1001'",
                "self",
                ["CUST-9527"],
            )
        self.assertEqual(cm.exception.error_code, "AUTH_403_NOT_OWNER")

    def test_insert_rejected(self):
        with self.assertRaises(SqlGuardError) as cm:
            validate("INSERT INTO core_customer VALUES (1)", "full")
        self.assertEqual(cm.exception.error_code, "SQL_NOT_SELECT")

    def test_advisor_out_of_scope_rejected(self):
        with self.assertRaises(SqlGuardError) as cm:
            validate(
                "SELECT * FROM core_holding WHERE customer_id = 'CUST-1004'",
                "assigned",
                ["CUST-1001"],
            )
        self.assertEqual(cm.exception.error_code, "AUTH_403_NOT_ASSIGNED")

    def test_ops_aggregate_allowed(self):
        r = validate("SELECT COUNT(DISTINCT customer_id) AS cnt FROM core_holding", "aggregate")
        self.assertTrue(r.allowed)

    def test_created_at_in_select_allowed(self):
        """Q17：列名 created_at 不应触发 create 关键字误杀。"""
        r = validate(
            "SELECT alert_id, created_at FROM jinrong_agent.risk_alert WHERE status='pending_review'",
            "full",
        )
        self.assertTrue(r.allowed)

    def test_gap_a_select_customer_id_column_without_where_denied_self(self):
        with self.assertRaises(SqlGuardError) as cm:
            validate(
                "SELECT customer_id, product_id, market_value FROM core_holding",
                "self",
                ["CUST-9527"],
            )
        self.assertEqual(cm.exception.error_code, "AUTH_403_SCOPE")

    def test_gap_a_select_customer_id_column_without_where_denied_assigned(self):
        with self.assertRaises(SqlGuardError) as cm:
            validate(
                "SELECT customer_id, product_id, market_value FROM core_holding",
                "assigned",
                ["CUST-9527"],
            )
        self.assertEqual(cm.exception.error_code, "AUTH_403_SCOPE")

    def test_gap_b_ops_risk_alert_denied(self):
        with self.assertRaises(SqlGuardError) as cm:
            validate("SELECT * FROM jinrong_agent.risk_alert", "aggregate")
        self.assertEqual(cm.exception.error_code, "AUTH_403_SCOPE")

    def test_gap_b_advisor_risk_alert_without_filter_denied(self):
        with self.assertRaises(SqlGuardError) as cm:
            validate("SELECT * FROM jinrong_agent.risk_alert", "assigned", ["CUST-9527"])
        self.assertEqual(cm.exception.error_code, "AUTH_403_SCOPE")

    def test_create_table_still_rejected(self):
        with self.assertRaises(SqlGuardError) as cm:
            validate("CREATE TABLE evil (id INT)", "full")
        self.assertIn(cm.exception.error_code, ("SQL_FORBIDDEN", "SQL_NOT_SELECT"))

    def test_inject_ownership(self):
        out = inject_ownership("SELECT * FROM core_holding", ["CUST-1", "CUST-2"])
        self.assertIn("CUST-1", out)

    def test_extract_tables(self):
        self.assertEqual(
            extract_tables("SELECT * FROM core_customer c JOIN core_holding h ON h.customer_id=c.customer_id"),
            ["core_customer", "core_holding"],
        )