Spaces:
Sleeping
Sleeping
| import sys | |
| sys.path.insert(0, ".") | |
| import traceback | |
| import yaml | |
| from models import Action, Observation, State | |
| from tasks import TASKS, TASK_GROUPS, get_task | |
| from graders import grade_submission | |
| from server.sql_environment import SQLDebugEnvironment | |
| def check(name: str, fn) -> bool: | |
| try: | |
| result = fn() | |
| if result is True or result is None: | |
| print(f" PASS: {name}") | |
| return True | |
| else: | |
| print(f" FAIL: {name}: {result}") | |
| return False | |
| except Exception as e: | |
| print(f" FAIL: {name}: {e}") | |
| traceback.print_exc() | |
| return False | |
| def test_models_importable(): | |
| a = Action(query="SELECT 1", explanation="test") | |
| assert a.query == "SELECT 1" | |
| o = Observation( | |
| schema_sql="CREATE TABLE t (id INTEGER);", | |
| sample_data="INSERT INTO t VALUES (1);", | |
| buggy_query="SELECT * t", | |
| task_description="test", | |
| expected_columns=["id"], | |
| ) | |
| s = State(episode_id="test-123") | |
| return True | |
| def test_all_tasks_defined(): | |
| assert len(TASKS) >= 9, f"Need 9+ tasks, got {len(TASKS)}" | |
| assert len(TASK_GROUPS) >= 3, f"Need 3 task groups, got {len(TASK_GROUPS)}" | |
| for gid, task_ids in TASK_GROUPS.items(): | |
| assert len(task_ids) >= 1, f"Group {gid} has no tasks" | |
| for tid in task_ids: | |
| assert tid in TASKS, f"Task {tid} in group {gid} not found in TASKS" | |
| return True | |
| def test_graders_correct_solutions(): | |
| failures = [] | |
| for task_id, task in TASKS.items(): | |
| score, feedback, _ = grade_submission(task_id, task["correct_query"]) | |
| if score != 1.0: | |
| failures.append(f"{task_id}: score={score} ({feedback[:60]})") | |
| if failures: | |
| return f"Graders failed on correct solutions: {failures}" | |
| return True | |
| def test_graders_wrong_query(): | |
| for task_id in list(TASKS.keys())[:3]: | |
| score, _, _ = grade_submission(task_id, "SELECT 'wrong'") | |
| assert score < 1.0, f"Wrong query got score={score} on {task_id}" | |
| return True | |
| def test_grader_score_range(): | |
| for task_id, task in TASKS.items(): | |
| s1, _, _ = grade_submission(task_id, task["correct_query"]) | |
| assert 0.0 <= s1 <= 1.0, f"{task_id} correct_query score out of range: {s1}" | |
| s2, _, _ = grade_submission(task_id, task["buggy_query"]) | |
| assert 0.0 <= s2 <= 1.0, f"{task_id} buggy_query score out of range: {s2}" | |
| assert s2 < 1.0, f"{task_id} buggy query should score < 1.0, got {s2}" | |
| return True | |
| def test_environment_reset_step_state(): | |
| env = SQLDebugEnvironment() | |
| for group in ["easy_syntax_fix", "medium_logic_fix", "hard_optimization_fix"]: | |
| obs = env.reset(task_group=group, seed=0) | |
| assert obs.schema_sql, "schema_sql is empty" | |
| assert obs.buggy_query, "buggy_query is empty" | |
| assert obs.task_description, "task_description is empty" | |
| assert obs.attempts_remaining == 5 | |
| assert not obs.done | |
| state = env.state() | |
| assert state.step_count == 0 | |
| assert state.task_id != "" | |
| action = Action(query="SELECT 1", explanation="test") | |
| obs2 = env.step(action) | |
| assert obs2.attempts_remaining == 4 | |
| assert obs2.reward is not None | |
| state2 = env.state() | |
| assert state2.step_count == 1 | |
| return True | |
| def test_reward_on_correct_answer(): | |
| env = SQLDebugEnvironment() | |
| obs = env.reset(task_id="easy_missing_from", seed=42) | |
| task = get_task(env._current_task_id) | |
| action = Action(query=task["correct_query"]) | |
| obs2 = env.step(action) | |
| assert ( | |
| obs2.reward >= 1.0 | |
| ), f"Correct answer should get reward >= 1.0, got {obs2.reward}" | |
| assert obs2.done, "Episode should be done after correct answer" | |
| assert obs2.success, "success should be True" | |
| return True | |
| def test_reward_partial_progress(): | |
| env = SQLDebugEnvironment() | |
| env.reset(task_id="easy_missing_from", seed=42) | |
| action = Action( | |
| query="SELECT first_name, last_name FROM employees ORDER BY last_name" | |
| ) | |
| obs2 = env.step(action) | |
| assert obs2.reward >= 0.0, "Should not penalize on first attempt" | |
| assert not obs2.success | |
| return True | |
| def test_max_steps_terminates(): | |
| env = SQLDebugEnvironment() | |
| env.reset(task_id="easy_missing_from", seed=42) | |
| for i in range(5): | |
| obs = env.step(Action(query="SELECT 1")) | |
| if obs.done: | |
| break | |
| assert obs.done, "Episode must terminate after max_steps" | |
| return True | |
| def test_openenv_yaml_valid(): | |
| with open("openenv.yaml", "r") as f: | |
| data = yaml.safe_load(f) | |
| required_keys = [ | |
| "name", | |
| "version", | |
| "description", | |
| "tasks", | |
| "action_space", | |
| "observation_space", | |
| ] | |
| for key in required_keys: | |
| assert key in data, f"openenv.yaml missing key: {key}" | |
| assert ( | |
| len(data["tasks"]) >= 3 | |
| ), f"Need 3+ tasks in openenv.yaml, got {len(data['tasks'])}" | |
| return True | |
| def test_deterministic_grading(): | |
| """Same query must return same score every time.""" | |
| query = "SELECT first_name, last_name FROM employees WHERE department_id = 10 ORDER BY last_name" | |
| scores = [grade_submission("easy_missing_from", query)[0] for _ in range(5)] | |
| assert len(set(scores)) == 1, f"Grading is not deterministic: {scores}" | |
| return True | |
| def main(): | |
| print("\n" + "=" * 60) | |
| print("SQL Debug Environment — Pre-Submission Validation") | |
| print("=" * 60 + "\n") | |
| checks = [ | |
| ("Models are importable (Action, Observation, State)", test_models_importable), | |
| ("All 9+ tasks defined in 3 task groups", test_all_tasks_defined), | |
| ("Graders return 1.0 on correct solutions", test_graders_correct_solutions), | |
| ("Graders return <1.0 on wrong queries", test_graders_wrong_query), | |
| ("All grader scores in range 0.0–1.0", test_grader_score_range), | |
| ("Environment reset/step/state cycle works", test_environment_reset_step_state), | |
| ("Correct answer gets reward >= 1.0", test_reward_on_correct_answer), | |
| ("Partial progress gets partial reward", test_reward_partial_progress), | |
| ("Episode terminates at max_steps", test_max_steps_terminates), | |
| ("openenv.yaml is valid", test_openenv_yaml_valid), | |
| ( | |
| "Grading is deterministic (same query = same score)", | |
| test_deterministic_grading, | |
| ), | |
| ] | |
| passed = 0 | |
| failed = 0 | |
| for name, fn in checks: | |
| ok = check(name, fn) | |
| if ok: | |
| passed += 1 | |
| else: | |
| failed += 1 | |
| print(f"\n{'='*60}") | |
| print(f"Results: {passed}/{len(checks)} checks passed") | |
| if failed == 0: | |
| print("ALL CHECKS PASSED — Ready to submit!") | |
| else: | |
| print(f"{failed} check(s) FAILED — Fix before submitting.") | |
| return 0 if failed == 0 else 1 | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |