Blame for validate.py

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())