Spaces:
Sleeping
Sleeping
File size: 2,808 Bytes
be62603 | 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 | """
Tests for the SQL Debug Environment lifecycle.
Run with: pytest tests/test_environment.py -v
"""
import sys
sys.path.insert(0, ".")
import pytest
from models import Action
from server.sql_environment import SQLDebugEnvironment
from tasks import TASKS, TASK_GROUPS, get_task
@pytest.fixture
def env():
return SQLDebugEnvironment()
def test_reset_returns_observation(env):
obs = env.reset()
assert obs.schema_sql
assert obs.buggy_query
assert obs.task_description
assert obs.attempts_remaining == 5
assert not obs.done
assert not obs.success
def test_reset_all_task_groups(env):
for group in TASK_GROUPS:
obs = env.reset(task_group=group)
assert obs.schema_sql
assert obs.buggy_query
def test_reset_specific_task(env):
obs = env.reset(task_id="easy_missing_from")
assert obs.attempts_remaining == 5
state = env.state()
assert state.task_id == "easy_missing_from"
def test_step_decrements_attempts(env):
env.reset(task_id="easy_missing_from")
obs = env.step(Action(query="SELECT 1"))
assert obs.attempts_remaining == 4
def test_correct_answer_succeeds(env):
env.reset(task_id="easy_missing_from")
task = get_task("easy_missing_from")
obs = env.step(Action(query=task["correct_query"]))
assert obs.success
assert obs.done
assert obs.reward >= 1.0
def test_max_steps_terminates(env):
env.reset(task_id="easy_missing_from")
for _ in range(5):
obs = env.step(Action(query="SELECT 1"))
if obs.done:
break
assert obs.done
def test_step_after_done_is_safe(env):
env.reset(task_id="easy_missing_from")
task = get_task("easy_missing_from")
env.step(Action(query=task["correct_query"]))
obs = env.step(Action(query="SELECT 1")) # after done
assert obs.done # should still be done
def test_state_step_count(env):
env.reset()
assert env.state().step_count == 0
env.step(Action(query="SELECT 1"))
assert env.state().step_count == 1
env.step(Action(query="SELECT 1"))
assert env.state().step_count == 2
def test_seed_reproducibility(env):
obs1 = env.reset(task_group="easy_syntax_fix", seed=42)
obs2 = env.reset(task_group="easy_syntax_fix", seed=42)
assert obs1.buggy_query == obs2.buggy_query
def test_partial_reward_on_progress(env):
env.reset(task_id="easy_missing_from")
obs = env.step(Action(
query="SELECT first_name, last_name FROM employees ORDER BY last_name"
))
assert obs.reward >= 0.0
assert not obs.success
def test_cumulative_reward_accumulates(env):
env.reset(task_id="easy_missing_from")
env.step(Action(query="SELECT 1"))
env.step(Action(query="SELECT 2"))
state = env.state()
assert state.cumulative_reward is not None
|