zhanghongyu_0626
a0f550646e
- Introduced a new `/analyze` endpoint in the analyst API to process analysis requests, allowing users to receive textual interpretations and chart specifications based on provided prompts. - Enhanced `analyst_schemas.py` with `AnalyzeRequest` and `ChartSpec` models to structure analysis requests and validate chart specifications. - Implemented chart validation logic in a new `analyst_chart.py` service, ensuring that chart types and fields are correctly specified and conform to allowed values. - Updated `AnalystAgent` to handle analysis requests, integrating the new logic for generating responses based on user prompts and data availability. - Added unit tests to verify the functionality of the new endpoint and validation mechanisms, ensuring robustness and reliability. This update significantly enhances the analytical capabilities of the application, providing users with improved tools for data interpretation and visualization.
117 lines
4.7 KiB
Python
117 lines
4.7 KiB
Python
"""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"],
|
|
)
|