Spaces:
Sleeping
Sleeping
| """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): | |
| 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() | |