Spaces:
Sleeping
Sleeping
File size: 3,189 Bytes
b634764 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """Guardrail + chart-shaping checks for the text-to-SQL agent.
Runnable: `python -m unittest tests.test_sql_agent`. Needs Postgres for the
execution tests; no LLM calls.
"""
import unittest
from sqlalchemy.exc import SQLAlchemyError
from backend.sql_agent import (
MAX_ROWS,
SqlAgentError,
build_chart,
run_sql,
sanitize_sql,
)
class SanitizeSqlTests(unittest.TestCase):
def test_adds_limit_and_strips_fence(self):
sql = sanitize_sql("```sql\nSELECT carrier FROM orders;\n```")
self.assertEqual(sql, f"SELECT carrier FROM orders LIMIT {MAX_ROWS}")
def test_keeps_existing_limit(self):
self.assertEqual(
sanitize_sql("SELECT carrier FROM orders LIMIT 5"),
"SELECT carrier FROM orders LIMIT 5",
)
def test_allows_leading_cte(self):
self.assertTrue(
sanitize_sql("WITH t AS (SELECT 1 AS v) SELECT v FROM t").startswith("WITH")
)
def test_rejects_writes_and_stacking(self):
for bad in (
"DELETE FROM orders",
"UPDATE orders SET status = 'delivered'",
"DROP TABLE orders",
"SELECT 1; SELECT 2",
"SELECT pg_sleep(10)",
"INSERT INTO orders (order_id) VALUES ('x')",
"",
):
with self.subTest(sql=bad), self.assertRaises(SqlAgentError):
sanitize_sql(bad)
class BuildChartTests(unittest.TestCase):
def test_bar_for_categories(self):
chart = build_chart(["label", "value"], [["UPS", 10], ["DHL", 4]])
self.assertEqual(chart["type"], "bar")
self.assertEqual(chart["labels"], ["UPS", "DHL"])
self.assertEqual(chart["values"], [10, 4])
def test_line_for_months(self):
chart = build_chart(["label", "value"], [["2025-01", 3], ["2025-02", 7]])
self.assertEqual(chart["type"], "line")
def test_none_when_not_chartable(self):
self.assertIsNone(build_chart(["value"], [[1], [2]])) # single column
self.assertIsNone(build_chart(["label", "value"], [["UPS", 10]])) # one row
class RunSqlTests(unittest.TestCase):
@classmethod
def setUpClass(cls):
try:
cls.total = run_sql("SELECT COUNT(*) AS n FROM orders")["rows"][0][0]
except SQLAlchemyError as error:
raise unittest.SkipTest(f"Postgres unavailable: {error}") from error
if not cls.total:
raise unittest.SkipTest("orders is empty; run `python -m backend.seed`")
def test_select_returns_columns_and_jsonable_rows(self):
result = run_sql("SELECT order_id, order_date, order_value_usd FROM orders LIMIT 1")
self.assertEqual(result["columns"], ["order_id", "order_date", "order_value_usd"])
order_id, order_date, value = result["rows"][0]
self.assertIsInstance(order_id, str)
self.assertRegex(order_date, r"^\d{4}-\d{2}-\d{2}$") # date -> ISO string
self.assertIsInstance(value, float) # Decimal -> float
def test_transaction_is_read_only(self):
with self.assertRaises(Exception):
run_sql("DELETE FROM orders WHERE FALSE")
if __name__ == "__main__":
unittest.main()
|