Spaces:
Sleeping
Sleeping
File size: 5,112 Bytes
822ca51 | 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 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | """Unit tests for server.safety -- pre/post-execution checks."""
import sqlite3
import pytest
from server.safety import (
SafetyViolation,
budget_cost,
check_post_execution,
check_pre_execution,
classify_sql,
)
CORE_TABLES = ["users", "products", "vendors"]
# classify_sql
class TestClassifySql:
def test_select(self):
assert classify_sql("SELECT * FROM users") == "SELECT"
def test_select_cte(self):
assert classify_sql("WITH cte AS (SELECT 1) SELECT * FROM cte") == "SELECT"
def test_pragma(self):
assert classify_sql("PRAGMA table_info(users)") == "META"
def test_insert(self):
assert classify_sql("INSERT INTO users VALUES (1, 'a@b.com')") == "MUTATION"
def test_update(self):
assert classify_sql("UPDATE users SET status='inactive' WHERE id=1") == "MUTATION"
def test_delete(self):
assert classify_sql("DELETE FROM users WHERE id=9") == "MUTATION"
def test_alter(self):
assert classify_sql("ALTER TABLE users ADD COLUMN phone TEXT") == "MUTATION"
def test_create(self):
assert classify_sql("CREATE TABLE foo (id INTEGER PRIMARY KEY)") == "MUTATION"
def test_leading_whitespace(self):
assert classify_sql(" \n SELECT id FROM users") == "SELECT"
# budget_cost
class TestBudgetCost:
def test_select_costs_1(self):
assert budget_cost("SELECT") == 1
def test_meta_costs_1(self):
assert budget_cost("META") == 1
def test_mutation_costs_5(self):
assert budget_cost("MUTATION") == 5
# check_pre_execution
class TestPreExecution:
def test_safe_select_returns_none(self):
assert check_pre_execution("SELECT * FROM users", CORE_TABLES) is None
def test_drop_core_table_triggers_violation(self):
v = check_pre_execution("DROP TABLE users", CORE_TABLES)
assert v is not None
assert v.kind == "SCHEMA_DESTRUCTION"
assert "users" in v.message
def test_drop_non_core_table_is_allowed(self):
v = check_pre_execution("DROP TABLE temp_backup", CORE_TABLES)
assert v is None
def test_drop_table_if_exists_core(self):
v = check_pre_execution("DROP TABLE IF EXISTS products", CORE_TABLES)
assert v is not None
assert v.kind == "SCHEMA_DESTRUCTION"
def test_delete_with_where_is_safe(self):
v = check_pre_execution(
"DELETE FROM users WHERE id = 9", CORE_TABLES
)
assert v is None
def test_delete_without_where_triggers_violation(self):
v = check_pre_execution("DELETE FROM users", CORE_TABLES)
assert v is not None
assert v.kind == "TRUNCATION"
def test_delete_from_with_newline_no_where(self):
v = check_pre_execution("DELETE\nFROM\nusers", CORE_TABLES)
assert v is not None
assert v.kind == "TRUNCATION"
def test_delete_with_subquery_where(self):
sql = "DELETE FROM users WHERE id IN (SELECT id FROM temp_dups)"
v = check_pre_execution(sql, CORE_TABLES)
assert v is None
def test_drop_core_table_case_insensitive(self):
v = check_pre_execution("drop table USERS", CORE_TABLES)
assert v is not None
assert v.kind == "SCHEMA_DESTRUCTION"
def test_pragma_is_safe(self):
assert check_pre_execution("PRAGMA foreign_key_check", CORE_TABLES) is None
def test_insert_is_safe(self):
assert (
check_pre_execution(
"INSERT INTO users VALUES (99, 'x@y.com', 'x', '2024-01-01', 'active')",
CORE_TABLES,
)
is None
)
# check_post_execution
class TestPostExecution:
def _make_conn_with_users(self, n: int) -> sqlite3.Connection:
conn = sqlite3.connect(":memory:")
conn.execute("CREATE TABLE users (id INTEGER PRIMARY KEY, email TEXT)")
for i in range(n):
conn.execute("INSERT INTO users VALUES (?, ?)", (i + 1, f"u{i}@test.com"))
conn.commit()
return conn
def test_no_loss_returns_none(self):
conn = self._make_conn_with_users(10)
initial = {"users": 10}
assert check_post_execution(conn, initial) is None
def test_small_loss_is_ok(self):
conn = self._make_conn_with_users(10)
conn.execute("DELETE FROM users WHERE id = 10")
conn.commit()
# floor = 10 (what we must protect); 9 >= 10*0.9 = 9 - OK
floor = {"users": 10}
assert check_post_execution(conn, floor) is None
def test_catastrophic_loss_triggers_violation(self):
conn = self._make_conn_with_users(10)
conn.execute("DELETE FROM users WHERE id > 2") # 8 deleted -> 2 remain (20%)
conn.commit()
floor = {"users": 10}
v = check_post_execution(conn, floor)
assert v is not None
assert v.kind == "CATASTROPHIC_LOSS"
assert "users" in v.message
def test_zero_initial_is_ignored(self):
conn = self._make_conn_with_users(0)
initial = {"users": 0}
assert check_post_execution(conn, initial) is None
|