Spaces:
Runtime error
Runtime error
File size: 6,841 Bytes
0cac9cf | 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | """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()
|