zhanghongyu_0626
bcd6175d4e
- Introduced `analyst_auth_adapter.py` for managing authentication context and access control for the data analysis agent. - Added new API endpoints in `analyst.py` for chat, dashboard, asset management, and metrics, utilizing the new authentication context. - Created Pydantic models in `analyst_schemas.py` for request and response structures, ensuring consistent data handling. - Updated SQL guard logic in `sql_guard.py` to enforce access restrictions based on user roles and contexts. - Implemented migration scripts for new database tables related to the data analysis agent, enhancing data management capabilities. - Removed legacy authentication code from `auth.py`, streamlining the authentication process. This update significantly enhances the data analysis capabilities, providing a robust framework for querying and managing data securely.
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"],
|
|
)
|