Spaces:
Runtime error
Runtime error
| """End-to-end tests for the async QueryEngine with a fake LLM. | |
| Author: mohamedgamal04 | |
| """ | |
| import json | |
| import unittest | |
| from pathlib import Path | |
| from tempfile import TemporaryDirectory | |
| from unittest.mock import patch | |
| import pandas as pd | |
| from queryquest.core.engine import QueryEngine | |
| from queryquest.core.models import EngineConfig | |
| from queryquest.core.policy import AutoApprovePolicy, DenyAllPolicy, Policy | |
| def _config(excel_dir: Path) -> EngineConfig: | |
| return EngineConfig( | |
| base_url="http://fake", | |
| api_key="key", | |
| model="fake-model", | |
| provider_name="fake", | |
| system_prompt="system", | |
| excel_dir=excel_dir, | |
| ) | |
| def _fake_llm(statements: list[str], explanation: str = "ok"): | |
| """Build an async stand-in for core.llm.generate_sql.""" | |
| async def _generate(config, user_prompt): | |
| raw = json.dumps({"sql_statements": statements, "explanation": explanation}) | |
| return raw, statements, explanation, None | |
| return _generate | |
| class _WritebackDenyPolicy(Policy): | |
| """Approves execution but refuses every write-back.""" | |
| async def approve_execution(self, statements) -> bool: | |
| return True | |
| async def approve_writeback(self, target) -> bool: | |
| return False | |
| class CoreEngineTests(unittest.IsolatedAsyncioTestCase): | |
| def _make_workbook(self, directory: Path) -> Path: | |
| path = directory / "listings.xlsx" | |
| pd.DataFrame({"id": [1, 2], "price": [10, 20]}).to_excel(path, index=False) | |
| return path | |
| async def test_select_returns_preview_rows(self) -> None: | |
| with TemporaryDirectory() as raw_dir: | |
| directory = Path(raw_dir) | |
| self._make_workbook(directory) | |
| engine = QueryEngine(_config(directory)) | |
| with patch("queryquest.core.engine.generate_sql", new=_fake_llm(["SELECT * FROM listings"])): | |
| result = await engine.run("show listings", AutoApprovePolicy(directory)) | |
| self.assertTrue(result.executed) | |
| self.assertIsNone(result.error) | |
| self.assertEqual(len(result.statements), 1) | |
| statement = result.statements[0] | |
| self.assertIsNone(statement.error) | |
| self.assertEqual(statement.row_count, 2) | |
| self.assertIn("price", statement.columns) | |
| async def test_update_writes_back_when_approved(self) -> None: | |
| with TemporaryDirectory() as raw_dir: | |
| directory = Path(raw_dir) | |
| path = self._make_workbook(directory) | |
| engine = QueryEngine(_config(directory)) | |
| sql = "UPDATE listings SET price = 99 WHERE id = 1" | |
| with patch("queryquest.core.engine.generate_sql", new=_fake_llm([sql])): | |
| result = await engine.run("bump price", AutoApprovePolicy(directory)) | |
| self.assertTrue(result.executed) | |
| self.assertTrue(result.wrote_back) | |
| saved = pd.read_excel(path) | |
| self.assertEqual(int(saved.loc[saved["id"] == 1, "price"].iloc[0]), 99) | |
| # The write-back target carries a preview of the rows about to change. | |
| target = result.writeback_targets[0] | |
| self.assertEqual(target.affected_rows, 1) | |
| self.assertEqual(len(target.preview_rows), 1) | |
| self.assertIn("price", target.preview_columns) | |
| async def test_zero_row_update_offers_no_writeback(self) -> None: | |
| with TemporaryDirectory() as raw_dir: | |
| directory = Path(raw_dir) | |
| path = self._make_workbook(directory) | |
| engine = QueryEngine(_config(directory)) | |
| sql = "UPDATE listings SET price = 5 WHERE id = 999" # matches nothing | |
| with patch("queryquest.core.engine.generate_sql", new=_fake_llm([sql])): | |
| result = await engine.run("noop", AutoApprovePolicy(directory)) | |
| self.assertTrue(result.executed) | |
| self.assertEqual(result.statements[0].row_count, 0) | |
| self.assertEqual(result.writeback_targets, []) | |
| self.assertFalse(result.wrote_back) | |
| saved = pd.read_excel(path) | |
| self.assertEqual(sorted(saved["price"].tolist()), [10, 20]) | |
| async def test_insert_reports_affected_count_and_saves(self) -> None: | |
| with TemporaryDirectory() as raw_dir: | |
| directory = Path(raw_dir) | |
| path = self._make_workbook(directory) | |
| engine = QueryEngine(_config(directory)) | |
| sql = "INSERT INTO listings VALUES (3, 30)" | |
| with patch("queryquest.core.engine.generate_sql", new=_fake_llm([sql])): | |
| result = await engine.run("add a row", AutoApprovePolicy(directory)) | |
| self.assertEqual(result.statements[0].row_count, 1) # not 0 despite DuckDB rowcount -1 | |
| self.assertTrue(result.wrote_back) | |
| saved = pd.read_excel(path) | |
| self.assertEqual(len(saved), 3) | |
| self.assertIn(3, saved["id"].tolist()) | |
| async def test_deny_all_skips_execution(self) -> None: | |
| with TemporaryDirectory() as raw_dir: | |
| directory = Path(raw_dir) | |
| self._make_workbook(directory) | |
| engine = QueryEngine(_config(directory)) | |
| with patch("queryquest.core.engine.generate_sql", new=_fake_llm(["SELECT * FROM listings"])): | |
| result = await engine.run("show listings", DenyAllPolicy()) | |
| self.assertFalse(result.executed) | |
| self.assertFalse(result.wrote_back) | |
| async def test_writeback_denied_leaves_file_unchanged(self) -> None: | |
| with TemporaryDirectory() as raw_dir: | |
| directory = Path(raw_dir) | |
| path = self._make_workbook(directory) | |
| engine = QueryEngine(_config(directory)) | |
| sql = "UPDATE listings SET price = 99 WHERE id = 1" | |
| with patch("queryquest.core.engine.generate_sql", new=_fake_llm([sql])): | |
| result = await engine.run("bump price", _WritebackDenyPolicy()) | |
| self.assertTrue(result.executed) | |
| self.assertFalse(result.wrote_back) | |
| saved = pd.read_excel(path) | |
| self.assertEqual(int(saved.loc[saved["id"] == 1, "price"].iloc[0]), 10) | |
| async def test_sandbox_blocks_file_read(self) -> None: | |
| with TemporaryDirectory() as raw_dir: | |
| directory = Path(raw_dir) | |
| self._make_workbook(directory) | |
| engine = QueryEngine(_config(directory)) | |
| with patch( | |
| "queryquest.core.engine.generate_sql", | |
| new=_fake_llm(["SELECT * FROM read_csv('/etc/passwd')"]), | |
| ): | |
| result = await engine.run("leak", AutoApprovePolicy(directory)) | |
| self.assertFalse(result.executed) | |
| self.assertEqual(len(result.statements), 1) | |
| self.assertIsNotNone(result.statements[0].error) | |
| if __name__ == "__main__": | |
| unittest.main() | |