ritvik360/nl2sql-bench
0
1"""2nl2sql-bench/tests/test_all.py3================================4Comprehensive test suite covering:5 - Database seeder (determinism + row counts)6 - Grader (all reward components, step penalty, edge cases)7 - Task registry (all 3 tasks load and produce valid examples)8 - Environment (reset, step, episode boundary, done logic)9 - Inference log format (regex checks on START / STEP / END)10 11Run with:12 pytest tests/ -v13or from project root:14 PYTHONPATH=.:server pytest tests/ -v15"""16 17from __future__ import annotations18 19import re20import sqlite321import sys22import os23from pathlib import Path24 25import pytest26 27# ── Path setup so tests can import from both project root and server/ ──────28ROOT = Path(__file__).parent.parent29SERVER = ROOT / "server"30sys.path.insert(0, str(ROOT))31sys.path.insert(0, str(SERVER))32 33# ── Fixtures ───────────────────────────────────────────────────────────────34 35@pytest.fixture(scope="session")36def db_conn():37 """Shared in-memory SQLite connection with full schema + seed data."""38 from db.seed import seed_database39 schema = (SERVER / "db" / "schema.sql").read_text()40 conn = sqlite3.connect(":memory:", check_same_thread=False)41 conn.row_factory = sqlite3.Row42 conn.executescript(schema)43 seed_database(conn)44 yield conn45 conn.close()46 47 48@pytest.fixture49def fresh_env():50 """A fresh NL2SQLEnvironment instance per test."""51 from environment import NL2SQLEnvironment52 return NL2SQLEnvironment()53 54 55# ══════════════════════════════════════════════════════════════════════════════56# 1. DATABASE SEEDER57# ══════════════════════════════════════════════════════════════════════════════58 59class TestSeeder:60 61 def test_categories_count(self, db_conn):62 row = db_conn.execute("SELECT COUNT(*) FROM categories").fetchone()63 assert row[0] == 8, "Should have exactly 8 categories"64 65 def test_products_count(self, db_conn):66 row = db_conn.execute("SELECT COUNT(*) FROM products").fetchone()67 assert row[0] == 64, "Should have 8 products × 8 categories = 64"68 69 def test_customers_count(self, db_conn):70 row = db_conn.execute("SELECT COUNT(*) FROM customers").fetchone()71 assert row[0] == 15072 73 def test_orders_exist(self, db_conn):74 row = db_conn.execute("SELECT COUNT(*) FROM orders").fetchone()75 assert row[0] > 100, "Should have a meaningful number of orders"76 77 def test_order_items_exist(self, db_conn):78 row = db_conn.execute("SELECT COUNT(*) FROM order_items").fetchone()79 assert row[0] > 20080 81 def test_reviews_exist(self, db_conn):82 row = db_conn.execute("SELECT COUNT(*) FROM reviews").fetchone()83 assert row[0] > 5084 85 def test_determinism(self, db_conn):86 """Seeding a second connection with the same seed gives identical counts."""87 from db.seed import seed_database88 schema = (SERVER / "db" / "schema.sql").read_text()89 conn2 = sqlite3.connect(":memory:")90 conn2.executescript(schema)91 seed_database(conn2)92 93 for tbl in ["categories", "products", "customers", "orders",94 "order_items", "reviews"]:95 c1 = db_conn.execute(f"SELECT COUNT(*) FROM {tbl}").fetchone()[0]96 c2 = conn2.execute(f"SELECT COUNT(*) FROM {tbl}").fetchone()[0]97 assert c1 == c2, f"Table {tbl} count mismatch: {c1} vs {c2}"98 conn2.close()99 100 def test_tiers_valid(self, db_conn):101 bad = db_conn.execute(102 "SELECT COUNT(*) FROM customers WHERE tier NOT IN ('bronze','silver','gold')"103 ).fetchone()[0]104 assert bad == 0105 106 def test_statuses_valid(self, db_conn):107 bad = db_conn.execute(108 "SELECT COUNT(*) FROM orders "109 "WHERE status NOT IN ('pending','processing','shipped','delivered','cancelled')"110 ).fetchone()[0]111 assert bad == 0112 113 def test_ratings_valid(self, db_conn):114 bad = db_conn.execute(115 "SELECT COUNT(*) FROM reviews WHERE rating < 1 OR rating > 5"116 ).fetchone()[0]117 assert bad == 0118 119 def test_referential_integrity(self, db_conn):120 """Order items should reference valid orders and products."""121 orphan_orders = db_conn.execute(122 "SELECT COUNT(*) FROM order_items oi "123 "LEFT JOIN orders o ON o.id = oi.order_id WHERE o.id IS NULL"124 ).fetchone()[0]125 assert orphan_orders == 0126 127 orphan_products = db_conn.execute(128 "SELECT COUNT(*) FROM order_items oi "129 "LEFT JOIN products p ON p.id = oi.product_id WHERE p.id IS NULL"130 ).fetchone()[0]131 assert orphan_products == 0132 133 134# ══════════════════════════════════════════════════════════════════════════════135# 2. GRADER136# ══════════════════════════════════════════════════════════════════════════════137 138class TestGrader:139 140 def test_exact_match_first_step(self):141 from grader import grade142 gt = [{"id": 1, "name": "Alice"}, {"id": 2, "name": "Bob"}]143 result = grade(144 actual_rows=gt.copy(),145 ground_truth_rows=gt,146 error=None,147 step=1,148 order_sensitive=False,149 )150 assert result.reward == pytest.approx(1.0)151 assert result.exact_match is True152 assert result.syntax_ok is True153 assert result.columns_match is True154 assert result.row_count_match is True155 assert result.step_penalty == 0.0156 157 def test_syntax_error_gives_zero(self):158 from grader import grade159 result = grade(160 actual_rows=None,161 ground_truth_rows=[{"x": 1}],162 error="near 'SELCT': syntax error",163 step=1,164 )165 assert result.reward == 0.0166 assert result.syntax_ok is False167 168 def test_step_penalty_applied(self):169 from grader import grade170 gt = [{"n": 1}]171 result = grade(172 actual_rows=gt.copy(),173 ground_truth_rows=gt,174 error=None,175 step=3, # penalty = (3-1)*0.05 = 0.10176 )177 assert result.reward == pytest.approx(1.0 - 0.10)178 assert result.step_penalty == pytest.approx(0.10)179 180 def test_columns_wrong_zero_higher_components(self):181 from grader import grade182 gt = [{"name": "Alice", "score": 10}]183 actual = [{"user": "Alice", "points": 10}] # wrong column names184 result = grade(actual_rows=actual, ground_truth_rows=gt, error=None, step=1)185 assert result.columns_match is False186 assert result.exact_match is False187 # Only syntax score: 0.10188 assert result.reward == pytest.approx(0.10)189 190 def test_correct_columns_wrong_rows(self):191 from grader import grade192 gt = [{"name": "Alice"}, {"name": "Bob"}]193 actual = [{"name": "Charlie"}, {"name": "Dave"}]194 result = grade(actual_rows=actual, ground_truth_rows=gt, error=None, step=1)195 assert result.columns_match is True196 assert result.row_count_match is True197 assert result.exact_match is False198 # syntax(0.10) + columns(0.20) + row_count(0.20) = 0.50199 assert result.reward == pytest.approx(0.50)200 201 def test_order_sensitive_wrong_order_is_not_exact(self):202 from grader import grade203 gt = [{"id": 1}, {"id": 2}]204 actual = [{"id": 2}, {"id": 1}] # reversed205 result = grade(206 actual_rows=actual,207 ground_truth_rows=gt,208 error=None,209 step=1,210 order_sensitive=True,211 )212 assert result.exact_match is False213 214 def test_order_insensitive_accepts_different_row_order(self):215 from grader import grade216 gt = [{"id": 1}, {"id": 2}]217 actual = [{"id": 2}, {"id": 1}] # different order but same content218 result = grade(219 actual_rows=actual,220 ground_truth_rows=gt,221 error=None,222 step=1,223 order_sensitive=False,224 )225 assert result.exact_match is True226 227 def test_penalty_never_makes_reward_negative(self):228 from grader import grade229 # Step 99 with syntax error → reward must be >= 0230 result = grade(231 actual_rows=None,232 ground_truth_rows=[{"x": 1}],233 error="some error",234 step=99,235 )236 assert result.reward >= 0.0237 238 def test_execute_query_blocks_writes(self, db_conn):239 from grader import execute_query240 rows, err = execute_query(db_conn, "INSERT INTO categories(name) VALUES ('x')")241 assert rows is None242 assert "not allowed" in err.lower() or "INSERT" in err243 244 def test_execute_query_returns_rows(self, db_conn):245 from grader import execute_query246 rows, err = execute_query(db_conn, "SELECT id, name FROM categories ORDER BY id")247 assert err is None248 assert len(rows) == 8249 assert "id" in rows[0]250 assert "name" in rows[0]251 252 def test_compute_ground_truth(self, db_conn):253 from grader import compute_ground_truth254 rows = compute_ground_truth(db_conn, "SELECT COUNT(*) AS n FROM customers")255 assert len(rows) == 1256 assert rows[0]["n"] == 150257 258 259# ══════════════════════════════════════════════════════════════════════════════260# 3. TASK REGISTRY261# ══════════════════════════════════════════════════════════════════════════════262 263class TestTasks:264 265 def test_all_tasks_registered(self):266 from tasks import all_task_names267 names = all_task_names()268 assert "simple-filter" in names269 assert "join-aggregation" in names270 assert "analytics-window" in names271 272 @pytest.mark.parametrize("task_name", [273 "simple-filter", "join-aggregation", "analytics-window"274 ])275 def test_task_has_examples(self, task_name):276 from tasks import get_task277 task = get_task(task_name)278 assert len(task.examples) >= 3, f"{task_name} needs at least 3 examples"279 280 @pytest.mark.parametrize("task_name", [281 "simple-filter", "join-aggregation", "analytics-window"282 ])283 def test_task_sql_runs_on_real_db(self, task_name, db_conn):284 """Every ground-truth SQL must execute cleanly against the seeded DB."""285 from tasks import get_task286 from grader import execute_query287 task = get_task(task_name)288 for ex in task.examples:289 rows, error = execute_query(db_conn, ex.sql)290 assert error is None, (291 f"Task {task_name!r} SQL failed:\n{ex.sql}\nError: {error}"292 )293 assert rows is not None294 295 @pytest.mark.parametrize("task_name", [296 "simple-filter", "join-aggregation", "analytics-window"297 ])298 def test_task_roundrobin(self, task_name):299 from tasks import get_task300 task = get_task(task_name)301 n = len(task.examples)302 seen = [task.next_example() for _ in range(n * 2)]303 # After n calls, second half should repeat first half304 assert seen[:n] == seen[n:]305 306 def test_schema_context_non_empty(self):307 from tasks import get_task308 task = get_task("simple-filter")309 ctx = task.schema_context()310 assert "customers" in ctx311 assert "orders" in ctx312 assert "products" in ctx313 314 315# ══════════════════════════════════════════════════════════════════════════════316# 4. ENVIRONMENT317# ══════════════════════════════════════════════════════════════════════════════318 319class TestEnvironment:320 321 def test_reset_returns_observation(self, fresh_env):322 obs = fresh_env.reset(task_name="simple-filter")323 assert obs.question != ""324 assert obs.schema_context != ""325 assert obs.task_name == "simple-filter"326 assert obs.done is False327 assert obs.step == 0328 assert obs.reward is None329 330 def test_reset_state(self, fresh_env):331 fresh_env.reset(task_name="join-aggregation")332 state = fresh_env.state333 assert state.task_name == "join-aggregation"334 assert state.task_difficulty == "medium"335 assert state.step_count == 0336 assert state.solved is False337 338 def test_step_increments_step_count(self, fresh_env):339 from models import NL2SQLAction340 fresh_env.reset(task_name="simple-filter")341 fresh_env.step(NL2SQLAction(query="SELECT 1"))342 assert fresh_env.state.step_count == 1343 344 def test_step_syntax_error_gives_nonzero_error(self, fresh_env):345 from models import NL2SQLAction346 fresh_env.reset(task_name="simple-filter")347 obs = fresh_env.step(NL2SQLAction(query="SELCT * FORM broken"))348 assert obs.last_error is not None349 assert obs.reward == 0.0350 351 def test_step_valid_query_returns_result(self, fresh_env):352 from models import NL2SQLAction353 fresh_env.reset(task_name="simple-filter")354 obs = fresh_env.step(NL2SQLAction(355 query="SELECT id, name FROM customers ORDER BY name LIMIT 5"356 ))357 assert obs.last_error is None358 assert len(obs.last_result) <= 5359 assert obs.reward >= 0.0360 361 def test_exact_match_ends_episode(self, fresh_env):362 """Submitting the exact ground-truth SQL should solve the episode."""363 from models import NL2SQLAction364 fresh_env.reset(task_name="simple-filter")365 # Get the ground truth SQL from the internal example366 gt_sql = fresh_env._example.sql367 obs = fresh_env.step(NL2SQLAction(query=gt_sql))368 assert obs.done is True369 assert fresh_env.state.solved is True370 assert obs.reward == pytest.approx(1.0) # step 1, full score371 372 def test_max_steps_ends_episode(self, fresh_env):373 """Exhausting all steps should end the episode even without solving."""374 from models import NL2SQLAction375 from environment import MAX_STEPS376 fresh_env.reset(task_name="analytics-window")377 obs = None378 for _ in range(MAX_STEPS):379 obs = fresh_env.step(NL2SQLAction(query="SELECT 1"))380 assert obs is not None381 assert obs.done is True382 383 def test_reset_clears_previous_episode(self, fresh_env):384 from models import NL2SQLAction385 fresh_env.reset(task_name="simple-filter")386 fresh_env.step(NL2SQLAction(query="SELECT 1"))387 # Second reset should clear state388 obs = fresh_env.reset(task_name="join-aggregation")389 assert fresh_env.state.step_count == 0390 assert obs.step == 0391 assert obs.task_name == "join-aggregation"392 393 @pytest.mark.parametrize("task_name", [394 "simple-filter", "join-aggregation", "analytics-window"395 ])396 def test_all_tasks_solvable(self, task_name):397 """Ground-truth SQL should always produce reward == 1.0 on step 1."""398 from environment import NL2SQLEnvironment399 from models import NL2SQLAction400 env = NL2SQLEnvironment()401 env.reset(task_name=task_name)402 gt_sql = env._example.sql403 obs = env.step(NL2SQLAction(query=gt_sql))404 assert obs.done is True405 assert obs.reward == pytest.approx(1.0), (406 f"Task {task_name!r}: ground-truth SQL did not score 1.0.\n"407 f"SQL: {gt_sql}\nError: {obs.last_error}\nReward: {obs.reward}"408 )409 410 def test_score_normalised_to_0_1(self, fresh_env):411 from models import NL2SQLAction412 fresh_env.reset(task_name="simple-filter")413 for _ in range(3):414 obs = fresh_env.step(NL2SQLAction(query="SELECT 1 AS x"))415 assert 0.0 <= obs.score <= 1.0416 417 def test_write_query_blocked(self, fresh_env):418 from models import NL2SQLAction419 fresh_env.reset(task_name="simple-filter")420 obs = fresh_env.step(NL2SQLAction(421 query="INSERT INTO categories(name) VALUES ('hack')"422 ))423 assert obs.last_error is not None424 assert "not allowed" in obs.last_error.lower() or "INSERT" in obs.last_error425 426 427# ══════════════════════════════════════════════════════════════════════════════428# 5. LOG FORMAT COMPLIANCE429# ══════════════════════════════════════════════════════════════════════════════430 431class TestLogFormat:432 """Validate that the inference.py log helpers emit correct format."""433 434 START_RE = re.compile(435 r"^\[START\] task=\S+ env=\S+ model=\S+$"436 )437 STEP_RE = re.compile(438 r"^\[STEP\] step=\d+ action=.+ reward=\d+\.\d{2} "439 r"done=(true|false) error=.+$"440 )441 END_RE = re.compile(442 r"^\[END\] success=(true|false) steps=\d+ score=\d+\.\d{3} "443 r"rewards=[\d.,]+$"444 )445 446 def _capture(self, func, *args, **kwargs) -> str:447 import io448 from contextlib import redirect_stdout449 buf = io.StringIO()450 with redirect_stdout(buf):451 func(*args, **kwargs)452 return buf.getvalue().strip()453 454 def test_log_start_format(self):455 sys.path.insert(0, str(ROOT))456 from inference import log_start457 out = self._capture(log_start, "simple-filter", "Qwen/Qwen2.5-72B")458 assert self.START_RE.match(out), f"Bad [START] format: {out!r}"459 460 def test_log_step_format_null_error(self):461 from inference import log_step462 out = self._capture(log_step, 1, "SELECT 1", 0.10, False, None)463 assert self.STEP_RE.match(out), f"Bad [STEP] format: {out!r}"464 465 def test_log_step_format_with_error(self):466 from inference import log_step467 out = self._capture(log_step, 2, "SELCT 1", 0.0, False, "syntax error")468 assert self.STEP_RE.match(out), f"Bad [STEP] format: {out!r}"469 470 def test_log_end_format_success(self):471 from inference import log_end472 out = self._capture(log_end, True, 3, 0.850, [0.50, 1.0, 1.0])473 assert self.END_RE.match(out), f"Bad [END] format: {out!r}"474 475 def test_log_end_format_failure(self):476 from inference import log_end477 out = self._capture(log_end, False, 5, 0.100, [0.1, 0.0, 0.0, 0.0, 0.0])478 assert self.END_RE.match(out), f"Bad [END] format: {out!r}"479 480 def test_reward_two_decimal_places(self):481 from inference import log_step482 out = self._capture(log_step, 1, "SELECT 1", 0.5, False, None)483 # reward= field must have exactly 2 decimal places484 match = re.search(r"reward=(\d+\.\d+)", out)485 assert match, "No reward= field found"486 assert len(match.group(1).split(".")[1]) == 2487 488 def test_score_three_decimal_places(self):489 from inference import log_end490 out = self._capture(log_end, True, 1, 1.0, [1.0])491 match = re.search(r"score=(\d+\.\d+)", out)492 assert match493 assert len(match.group(1).split(".")[1]) == 3494