Spaces:
Sleeping
Sleeping
File size: 3,943 Bytes
cf796c5 8d73fdc cf796c5 | 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 | import pytest
from src.nodes.sql_safety import check_sql_safety
from src.state import init_state
def test_sql_safety_allows_select():
"""Test safety checker allows valid SELECT query"""
state = init_state("test")
state["sql_query"] = "SELECT * FROM orders WHERE order_id = 'abc123'"
result = check_sql_safety(state)
assert result["sql_query"] is not None
# sql_result is only set when there's an error, so check it's None or has no error
sql_result = result.get("sql_result")
assert sql_result is None or (isinstance(sql_result, dict) and sql_result.get("error") is None)
def test_sql_safety_blocks_drop():
"""Test safety checker blocks DROP statement"""
state = init_state("test")
state["sql_query"] = "DROP TABLE orders"
result = check_sql_safety(state)
assert result["sql_query"] is None
assert result["sql_result"]["error"] == "unsafe_query"
def test_sql_safety_blocks_delete():
"""Test safety checker blocks DELETE statement"""
state = init_state("test")
state["sql_query"] = "DELETE FROM orders WHERE 1=1"
result = check_sql_safety(state)
assert result["sql_query"] is None
assert result["sql_result"]["error"] == "unsafe_query"
def test_sql_safety_blocks_insert():
"""Test safety checker blocks INSERT statement"""
state = init_state("test")
state["sql_query"] = "INSERT INTO orders VALUES (1, 2, 3)"
result = check_sql_safety(state)
assert result["sql_query"] is None
def test_sql_safety_blocks_update():
"""Test safety checker blocks UPDATE statement"""
state = init_state("test")
state["sql_query"] = "UPDATE orders SET status = 'cancelled'"
result = check_sql_safety(state)
assert result["sql_query"] is None
def test_sql_safety_blocks_alter():
"""Test safety checker blocks ALTER statement"""
state = init_state("test")
state["sql_query"] = "ALTER TABLE orders ADD COLUMN test TEXT"
result = check_sql_safety(state)
assert result["sql_query"] is None
def test_sql_safety_blocks_truncate():
"""Test safety checker blocks TRUNCATE statement"""
state = init_state("test")
state["sql_query"] = "TRUNCATE TABLE orders"
result = check_sql_safety(state)
assert result["sql_query"] is None
def test_sql_safety_blocks_sqlite_master():
"""Test safety checker blocks access to sqlite_master"""
state = init_state("test")
state["sql_query"] = "SELECT * FROM sqlite_master"
result = check_sql_safety(state)
assert result["sql_query"] is None
assert result["sql_result"]["error"] == "system_table_access"
def test_sql_safety_blocks_union_injection():
"""Test safety checker blocks UNION injection attempt"""
state = init_state("test")
state["sql_query"] = "SELECT * FROM orders UNION SELECT * FROM customers"
result = check_sql_safety(state)
assert result["sql_query"] is None
def test_sql_safety_blocks_exec():
"""Test safety checker blocks EXEC statement"""
state = init_state("test")
state["sql_query"] = "EXEC sp_executesql"
result = check_sql_safety(state)
assert result["sql_query"] is None
def test_sql_safety_blocks_non_select():
"""Test safety checker blocks queries not starting with SELECT"""
state = init_state("test")
state["sql_query"] = "SHOW TABLES"
result = check_sql_safety(state)
assert result["sql_query"] is None
assert result["sql_result"]["error"] == "non_select"
def test_sql_safety_case_insensitive():
"""Test safety checker is case insensitive"""
state = init_state("test")
state["sql_query"] = "drop table orders"
result = check_sql_safety(state)
assert result["sql_query"] is None
def test_sql_safety_handles_empty_query():
"""Test safety checker handles empty query gracefully"""
state = init_state("test")
state["sql_query"] = ""
result = check_sql_safety(state)
assert result["sql_result"]["error"] == "no_query"
|