56 lines
2.0 KiB
Python
56 lines
2.0 KiB
Python
|
"""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"],
|
||
|
|
)
|