Spaces:
Sleeping
Sleeping
| """Integration tests for DBEREnvironment -- step/reset/state API.""" | |
| import pytest | |
| from db_er.models import DBERAction, DBERObservation, DBERState | |
| from server.environment import DBEREnvironment | |
| # Helpers | |
| def make_action(action_type: str, **kwargs) -> DBERAction: | |
| return DBERAction(action_type=action_type, **kwargs) | |
| def sql(query: str) -> DBERAction: | |
| return make_action("execute_sql", query=query) | |
| def submit(notes: str = "done") -> DBERAction: | |
| return make_action("submit_resolution", notes=notes) | |
| # reset() | |
| class TestReset: | |
| def test_reset_returns_observation(self): | |
| env = DBEREnvironment() | |
| obs = env.reset(task_id=1) | |
| assert isinstance(obs, DBERObservation) | |
| assert obs.task_id == 1 | |
| assert obs.budget_remaining == 100 | |
| assert obs.step_count == 0 | |
| assert obs.done is False | |
| assert "users" in obs.incident_ticket.lower() or "UNIQUE" in obs.incident_ticket | |
| def test_reset_seed_picks_task(self): | |
| env = DBEREnvironment() | |
| obs = env.reset(seed=0) | |
| assert obs.task_id == 1 | |
| obs = env.reset(seed=1) | |
| assert obs.task_id == 2 | |
| obs = env.reset(seed=2) | |
| assert obs.task_id == 3 | |
| def test_reset_overrides_task_id(self): | |
| env = DBEREnvironment() | |
| obs = env.reset(seed=0, task_id=3) | |
| assert obs.task_id == 3 | |
| def test_reset_starts_with_violations(self): | |
| env = DBEREnvironment() | |
| obs1 = env.reset(task_id=1) | |
| assert obs1.violation_count > 0 # 3 duplicate rows | |
| def test_reset_gives_clean_state(self): | |
| env = DBEREnvironment() | |
| obs = env.reset(task_id=1) | |
| # Take a step | |
| env.step(sql("SELECT * FROM users")) | |
| # Reset again - must return to clean state | |
| obs2 = env.reset(task_id=1) | |
| assert obs2.step_count == 0 | |
| assert obs2.budget_remaining == 100 | |
| def test_double_reset_isolated(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| # Delete all users (should not persist after reset) | |
| env.step(sql("DELETE FROM users WHERE id >= 9")) | |
| obs2 = env.reset(task_id=1) | |
| assert obs2.violation_count == 3 # duplicates back | |
| def test_invalid_task_id_raises(self): | |
| env = DBEREnvironment() | |
| with pytest.raises(ValueError): | |
| env.reset(task_id=99) | |
| # step() - budget mechanics | |
| class TestBudget: | |
| def test_select_costs_1(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| obs = env.step(sql("SELECT * FROM users")) | |
| assert obs.budget_remaining == 99 | |
| def test_mutation_costs_5(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| obs = env.step(sql("DELETE FROM users WHERE id = 9")) | |
| assert obs.budget_remaining == 95 | |
| def test_pragma_costs_1(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| obs = env.step(sql("PRAGMA table_info(users)")) | |
| assert obs.budget_remaining == 99 | |
| def test_step_count_increments(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| for i in range(3): | |
| obs = env.step(sql("SELECT 1")) | |
| assert obs.step_count == 3 | |
| def test_budget_exhaustion_ends_episode(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| # Use up budget with mutations | |
| for _ in range(20): | |
| obs = env.step(sql("SELECT 1")) # 1 point each | |
| # Keep draining | |
| for _ in range(80): | |
| obs = env.step(sql("SELECT 1")) | |
| # At 0 budget, done should be True | |
| assert obs.done is True or obs.budget_remaining <= 0 | |
| # step() - SQL execution | |
| class TestSqlExecution: | |
| def test_select_returns_rows(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| obs = env.step(sql("SELECT id, email FROM users ORDER BY id")) | |
| assert isinstance(obs.last_query_result, list) | |
| assert len(obs.last_query_result) == 11 # 8 + 3 duplicates | |
| def test_mutation_returns_rows_affected(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| obs = env.step(sql("DELETE FROM users WHERE id = 9")) | |
| assert obs.last_action_type == "MUTATION" | |
| def test_bad_sql_returns_error(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| obs = env.step(sql("SELECT * FROM nonexistent_table_xyz")) | |
| assert obs.last_action_type == "ERROR" | |
| assert "error" in obs.error_logs.lower() or "no such table" in obs.error_logs.lower() | |
| def test_zero_row_delete_penalised(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| # Delete a row that doesn't exist | |
| obs = env.step(sql("DELETE FROM users WHERE id = 9999")) | |
| assert obs.reward is not None and obs.reward < 0 | |
| def test_spam_penalty_on_repeated_select(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| for _ in range(3): | |
| obs = env.step(sql("SELECT COUNT(*) FROM users")) | |
| # Third execution should incur spam penalty | |
| assert obs.reward is not None and obs.reward <= 0 | |
| # step() - intermediate rewards | |
| class TestIntermediateRewards: | |
| def test_violation_reduction_gives_positive_reward(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| # Delete 1 duplicate -> violation_count 3->2 -> +0.2 reward | |
| obs = env.step(sql("DELETE FROM users WHERE id = 9")) | |
| assert obs.reward is not None and obs.reward >= 0.2 | |
| def test_multiple_violation_reductions(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| # Delete all 3 duplicates at once -> violation_count 3->0 -> +0.6 | |
| obs = env.step(sql(""" | |
| DELETE FROM users | |
| WHERE id NOT IN (SELECT MIN(id) FROM users GROUP BY email) | |
| """)) | |
| assert obs.reward is not None and obs.reward >= 0.6 | |
| assert obs.violation_count == 0 | |
| # step() - safety violations | |
| class TestSafetyViolations: | |
| def test_drop_core_table_ends_episode(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| obs = env.step(sql("DROP TABLE users")) | |
| assert obs.done is True | |
| assert obs.reward == -1.0 | |
| def test_drop_non_core_table_is_ok(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=2) | |
| # Create and drop a temp table | |
| env.step(sql("CREATE TABLE temp_work (id INTEGER PRIMARY KEY)")) | |
| obs = env.step(sql("DROP TABLE temp_work")) | |
| assert obs.done is False | |
| assert obs.reward != -1.0 | |
| def test_bare_delete_ends_episode(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| obs = env.step(sql("DELETE FROM users")) | |
| assert obs.done is True | |
| assert obs.reward == -1.0 | |
| def test_catastrophic_row_loss_ends_episode(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| # Golden DB has 8 users. Safety floor = 90% of 8 = 7.2. | |
| # Delete 9 of 11 rows -> 2 remain, which is < 7.2 - catastrophic. | |
| obs = env.step(sql("DELETE FROM users WHERE id > 2")) | |
| assert obs.done is True | |
| assert obs.reward == -1.0 | |
| def test_safe_delete_with_where_does_not_end_episode(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| # Delete the 3 duplicate rows (11->8). Golden floor = 8 rows. | |
| # 8 >= 8*0.9 = 7.2 -> safe, no catastrophic loss. | |
| obs = env.step(sql("DELETE FROM users WHERE id IN (9, 10, 11)")) | |
| assert obs.done is False | |
| # submit_resolution() | |
| class TestSubmitResolution: | |
| def test_submit_without_fix_gives_low_score(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| obs = env.step(submit("Nothing fixed.")) | |
| assert obs.done is True | |
| # Grader score for unrepaired DB should be < 1.0 | |
| # Final reward = score - 0.5 penalty (if score < 0.5) | |
| assert obs.reward < 1.0 | |
| def test_submit_after_fix_gives_high_score(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| env.step(sql(""" | |
| DELETE FROM users | |
| WHERE id NOT IN (SELECT MIN(id) FROM users GROUP BY email) | |
| """)) | |
| obs = env.step(submit("Deleted duplicate rows.")) | |
| assert obs.done is True | |
| assert obs.reward is not None and obs.reward >= 0.9 | |
| def test_submit_ends_episode(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| obs = env.step(submit("done")) | |
| assert obs.done is True | |
| # Further steps should raise | |
| with pytest.raises(RuntimeError): | |
| env.step(sql("SELECT 1")) | |
| # state property | |
| class TestState: | |
| def test_state_reflects_current_episode(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=2) | |
| env.step(sql("SELECT * FROM vendors")) | |
| state = env.state | |
| assert isinstance(state, DBERState) | |
| assert state.task_id == 2 | |
| assert state.budget_remaining == 99 | |
| assert state.step_count == 1 | |
| def test_state_tracks_violations(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| state_before = env.state | |
| import json | |
| env.step(sql("DELETE FROM users WHERE id = 9")) | |
| state_after = env.state | |
| assert state_after.violation_count < state_before.violation_count | |
| # close() | |
| class TestClose: | |
| def test_close_is_idempotent(self): | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| env.close() | |
| env.close() # should not raise | |
| def test_step_after_close_still_works(self): | |
| # close() is a no-op (class-level shared state for HTTP compatibility), | |
| # so step() after close() should succeed, not raise. | |
| env = DBEREnvironment() | |
| env.reset(task_id=1) | |
| env.close() | |
| obs = env.step(sql("SELECT 1")) | |
| assert obs is not None | |