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.
52 lines
1.8 KiB
Python
52 lines
1.8 KiB
Python
"""guardrail 数字护栏单元测试。"""
|
|
import unittest
|
|
|
|
from app.model.analyst_schemas import TableData
|
|
from app.service.guardrail import check_numbers, extract_numbers, result_numbers, verify
|
|
|
|
|
|
class TestGuardrail(unittest.TestCase):
|
|
def _table(self):
|
|
return TableData(
|
|
columns=["risk_code", "cnt"],
|
|
rows=[["C1", 4], ["C2", 6]],
|
|
)
|
|
|
|
def test_extract_numbers(self):
|
|
self.assertEqual(extract_numbers("共 2 个，高风险 6 人，占比 5%"), [2.0, 6.0, 5.0])
|
|
|
|
def test_extract_comma_numbers(self):
|
|
self.assertEqual(extract_numbers("2,625,000.00 元 和 1,229,150 元"), [2625000.0, 1229150.0])
|
|
|
|
def test_correct_answer_no_issues(self):
|
|
self.assertEqual(check_numbers("共 2 个风险等级，高风险 6 人", self._table()), [])
|
|
|
|
def test_wrong_number_flagged(self):
|
|
# 表格只有 4/6/合计10/行数2，答案说 123 应被拦截
|
|
issues = check_numbers("金额加起来是 123 万元", self._table())
|
|
self.assertIn(123.0, issues)
|
|
|
|
def test_result_numbers(self):
|
|
nums = result_numbers(self._table())
|
|
self.assertIn(2.0, nums) # 行数
|
|
self.assertIn(4.0, nums)
|
|
self.assertIn(6.0, nums)
|
|
self.assertIn(10.0, nums) # 数值列求和
|
|
|
|
def test_verify_wrong_fails(self):
|
|
r = verify("共 999 个", self._table(), data_as_of="2026-09-04")
|
|
self.assertFalse(r.passed)
|
|
self.assertIn(999.0, r.issues)
|
|
|
|
def test_zero_valid(self):
|
|
t = TableData(columns=["c"], rows=[])
|
|
self.assertEqual(check_numbers("结果为 0", t), [])
|
|
|
|
def test_wan_scale_valid(self):
|
|
t = TableData(columns=["v"], rows=[[1234567]])
|
|
self.assertEqual(check_numbers("约 123 万元", t), [])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|