Hariprita/nl2sql-openenv
0
1"""2test_local.py — Integration tests for the SQL Agent Environment.3 4Starts the FastAPI server as a subprocess, exercises every endpoint,5validates rewards, and prints ALL TESTS PASSED on success.6"""7 8import json9import subprocess10import sys11import time12import os13 14# Force UTF-8 output on Windows so Unicode characters in feedback don't crash prints15if sys.platform == "win32":16 import io17 sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")18 sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8", errors="replace")19 20import requests21 22BASE_URL = "http://localhost:7860"23SERVER_PROC = None24 25 26# ---------------------------------------------------------------------------27# Server lifecycle28# ---------------------------------------------------------------------------29 30def start_server():31 global SERVER_PROC32 env = os.environ.copy()33 SERVER_PROC = subprocess.Popen(34 [sys.executable, "-m", "uvicorn", "server.app:app",35 "--host", "0.0.0.0", "--port", "7860"],36 stdout=subprocess.PIPE,37 stderr=subprocess.PIPE,38 env=env,39 )40 # Wait up to 30 s for the server to be ready41 for _ in range(30):42 try:43 r = requests.get(f"{BASE_URL}/health", timeout=2)44 if r.status_code == 200:45 print(f"[OK] Server ready (pid={SERVER_PROC.pid})")46 return47 except requests.exceptions.ConnectionError:48 pass49 time.sleep(1)50 stop_server()51 raise RuntimeError("Server did not start within 30 s.")52 53 54def stop_server():55 if SERVER_PROC is not None:56 SERVER_PROC.terminate()57 try:58 SERVER_PROC.wait(timeout=5)59 except subprocess.TimeoutExpired:60 SERVER_PROC.kill()61 print("[OK] Server stopped.")62 63 64# ---------------------------------------------------------------------------65# Helpers66# ---------------------------------------------------------------------------67 68def check(condition: bool, label: str):69 if condition:70 print(f" [PASS] {label}")71 else:72 print(f" [FAIL] {label}")73 stop_server()74 sys.exit(1)75 76 77# ---------------------------------------------------------------------------78# Tests79# ---------------------------------------------------------------------------80 81def test_health():82 print("\n--- /health ---")83 r = requests.get(f"{BASE_URL}/health", timeout=5)84 check(r.status_code == 200, "HTTP 200")85 data = r.json()86 check(data.get("status") == "ok", f"status==ok (got {data})")87 88 89def test_tasks():90 print("\n--- GET /tasks ---")91 r = requests.get(f"{BASE_URL}/tasks", timeout=5)92 check(r.status_code == 200, "HTTP 200")93 data = r.json()94 tasks = data.get("tasks", [])95 check(len(tasks) == 3, f"3 tasks defined (got {len(tasks)})")96 ids = [t["id"] for t in tasks]97 check("simple_select" in ids, "simple_select task present")98 check("join_aggregation" in ids, "join_aggregation task present")99 check("window_ranking" in ids, "window_ranking task present")100 101 102def test_reset():103 print("\n--- POST /reset ---")104 r = requests.post(f"{BASE_URL}/reset", json={}, timeout=5)105 check(r.status_code == 200, "HTTP 200")106 data = r.json()107 print(f" question : {data.get('question','')[:80]}")108 check("schema" in data, "response has 'schema'")109 check("question" in data, "response has 'question'")110 check("task_id" in data, "response has 'task_id'")111 check(data.get("done") is False, "done==False on reset")112 check(data.get("reward") == 0.0, "reward==0.0 on reset")113 return data114 115 116def test_step_correct(obs_data: dict):117 print("\n--- POST /step (correct query — task 1) ---")118 sql = (119 "SELECT name, city FROM customers "120 "WHERE country = 'United States' "121 "ORDER BY name"122 )123 r = requests.post(f"{BASE_URL}/step", json={"sql_query": sql}, timeout=5)124 check(r.status_code == 200, "HTTP 200")125 data = r.json()126 reward = data.get("reward", -1)127 print(f" reward : {reward}")128 print(f" feedback : {data.get('feedback','')[:100]}")129 print(f" result : {str(data.get('result',''))[:120]}")130 check(reward >= 0.7, f"reward >= 0.7 for correct US customers query (got {reward})")131 return data132 133 134def test_step_bad_query():135 print("\n--- POST /step (non-existent table — should give 0.0) ---")136 # Reset first so we are on task 1137 requests.post(f"{BASE_URL}/reset", json={}, timeout=5)138 r = requests.post(139 f"{BASE_URL}/step",140 json={"sql_query": "SELECT * FROM nonexistent_table"},141 timeout=5,142 )143 check(r.status_code == 200, "HTTP 200")144 data = r.json()145 reward = data.get("reward", -1)146 print(f" reward : {reward}")147 print(f" feedback : {data.get('feedback','')[:100]}")148 check(reward == 0.0, f"reward==0.0 for bad table name (got {reward})")149 150 151def test_step_forbidden():152 print("\n--- POST /step (DROP TABLE — should give 0.0) ---")153 requests.post(f"{BASE_URL}/reset", json={}, timeout=5)154 r = requests.post(155 f"{BASE_URL}/step",156 json={"sql_query": "DROP TABLE customers"},157 timeout=5,158 )159 check(r.status_code == 200, "HTTP 200")160 data = r.json()161 reward = data.get("reward", -1)162 print(f" reward : {reward}")163 check(reward == 0.0, f"reward==0.0 for DROP statement (got {reward})")164 165 166def test_state():167 print("\n--- GET /state ---")168 r = requests.get(f"{BASE_URL}/state", timeout=5)169 check(r.status_code == 200, "HTTP 200")170 data = r.json()171 print(f" state : {json.dumps(data)}")172 check("episode_id" in data, "state has episode_id")173 check("step_count" in data, "state has step_count")174 check("cumulative_reward" in data, "state has cumulative_reward")175 check("done" in data, "state has done")176 177 178def test_full_episode():179 print("\n--- Full episode walkthrough (3 tasks) ---")180 requests.post(f"{BASE_URL}/reset", json={}, timeout=5)181 182 queries = [183 # Task 1 — easy184 "SELECT name, city FROM customers WHERE country = 'United States' ORDER BY name",185 # Task 2 — medium186 (187 "SELECT p.category, SUM(oi.quantity * oi.unit_price) AS total_revenue "188 "FROM order_items oi "189 "JOIN products p ON oi.product_id = p.product_id "190 "GROUP BY p.category "191 "ORDER BY total_revenue DESC"192 ),193 # Task 3 — hard (window function with SQLite RANK())194 (195 "SELECT c.name, MAX(o.order_date) AS most_recent_order, "196 "RANK() OVER (ORDER BY SUM(o.total_amount) DESC) AS rank "197 "FROM customers c "198 "JOIN orders o ON c.customer_id = o.customer_id "199 "GROUP BY c.customer_id, c.name "200 "ORDER BY rank"201 ),202 ]203 204 total_reward = 0.0205 for i, sql in enumerate(queries, 1):206 r = requests.post(f"{BASE_URL}/step", json={"sql_query": sql}, timeout=5)207 check(r.status_code == 200, f"task {i}: HTTP 200")208 data = r.json()209 reward = data.get("reward", 0.0)210 total_reward += reward211 print(f" Task {i} reward: {reward:.2f} | feedback: {data.get('feedback','')[:80]}")212 213 check(total_reward > 0, f"episode produced non-zero reward (got {total_reward:.2f})")214 print(f" Total episode reward: {total_reward:.2f}")215 216 217# ---------------------------------------------------------------------------218# Entry point219# ---------------------------------------------------------------------------220 221if __name__ == "__main__":222 print("Starting SQL Agent Environment tests...")223 start_server()224 try:225 test_health()226 test_tasks()227 obs = test_reset()228 test_step_correct(obs)229 test_step_bad_query()230 test_step_forbidden()231 test_state()232 test_full_episode()233 finally:234 stop_server()235 236 print("\n" + "=" * 40)237 print("ALL TESTS PASSED")238 print("=" * 40)239