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

from app.service.sql_guard import SqlGuardError, extract_tables, inject_ownership, 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_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_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"],
        )