tests/unit/test_sql_safety.py

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"