tests/test_execution_guard.py

"""test_execution_guard.py — Security and safety tests.

Covers:
  - Malicious SQL injection attempts
  - Forbidden tokens (DROP, INSERT, DELETE, etc.)
  - Multi-statement attacks
  - Malformed/broken JSON
  - Empty/missing inputs
  - Non-SELECT queries
"""

import pytest
from data_engine import extract_sql, validate_sql, execute_safe


# ── SQL extraction ─────────────────────────────────────────────────────

class TestExtractSql:
    """extract_sql() should handle all valid input formats."""

    def test_extract_sql_block(self):
        assert extract_sql("```sql\nSELECT 1\n```") == "SELECT 1"

    def test_extract_generic_code_block(self):
        assert extract_sql("```\nSELECT * FROM x\n```") == "SELECT * FROM x"

    def test_extract_json_envelope(self):
        result = extract_sql('{"sql": "SELECT COUNT(*) FROM enrollment", "explanation": "test"}')
        assert result == "SELECT COUNT(*) FROM enrollment"

    def test_extract_json_embedded_in_text(self):
        result = extract_sql('Here: {"sql": "SELECT 1"} rest')
        assert result == "SELECT 1"

    def test_extract_fallback_raw(self):
        assert extract_sql("SELECT 42") == "SELECT 42"

    def test_extract_strips_semicolons(self):
        assert extract_sql("```sql\nSELECT 1;\n```") == "SELECT 1"

    def test_extract_handles_multiline_sql(self):
        result = extract_sql("```sql\nSELECT a,\nb,\nc\nFROM t\n```")
        assert "SELECT a," in result
        assert "FROM t" in result

    def test_extract_json_without_sql_key(self):
        """JSON without 'sql' key should fall through to raw."""
        result = extract_sql('{"wrong": "key"}')
        assert result == '{"wrong": "key"}'

    def test_extract_malformed_json(self):
        """Malformed JSON should fall through to ```sql``` or raw."""
        result = extract_sql('{"sql": "SELECT 1"')
        # Broken JSON (missing closing brace) — should fall through
        assert "SELECT 1" in result


# ── Static validation ──────────────────────────────────────────────────

class TestValidateSqlStatic:
    """validate_sql() static checks (no DB connection needed)."""

    def test_rejects_empty(self):
        with pytest.raises(ValueError, match="Empty"):
            validate_sql("")

    def test_rejects_no_select(self):
        with pytest.raises(ValueError, match="No SELECT"):
            validate_sql("42")

    def test_rejects_drop(self):
        with pytest.raises(ValueError, match="drop"):
            validate_sql("DROP TABLE enrollment")

    def test_rejects_insert(self):
        with pytest.raises(ValueError, match="insert"):
            validate_sql("INSERT INTO enrollment VALUES (1)")

    def test_rejects_delete(self):
        with pytest.raises(ValueError, match="delete"):
            validate_sql("DELETE FROM attendance")

    def test_rejects_update(self):
        with pytest.raises(ValueError, match="update"):
            validate_sql("UPDATE enrollment SET x = 1")

    def test_rejects_alter(self):
        with pytest.raises(ValueError, match="alter"):
            validate_sql("ALTER TABLE enrollment ADD COLUMN x INT")

    def test_rejects_truncate(self):
        with pytest.raises(ValueError, match="truncate"):
            validate_sql("TRUNCATE TABLE enrollment")

    def test_rejects_create(self):
        with pytest.raises(ValueError, match="create"):
            validate_sql("CREATE TABLE hack (x INT)")

    def test_rejects_pragma(self):
        with pytest.raises(ValueError, match="pragma"):
            validate_sql("PRAGMA database_list")

    def test_allows_valid_select(self):
        validate_sql("SELECT * FROM enrollment")


# ── Schema-aware validation ────────────────────────────────────────────

class TestValidateSqlSchema:
    """validate_sql() schema-aware checks (with DB connection)."""

    def test_rejects_unknown_column(self, db):
        with pytest.raises(ValueError) as exc:
            validate_sql("SELECT fake_column FROM enrollment", conn=db)
        assert "fake_column" in str(exc.value).lower()

    def test_rejects_unknown_table(self, db):
        with pytest.raises(ValueError, match="exist"):
            validate_sql("SELECT * FROM nonexistent_table", conn=db)

    def test_accepts_valid_query(self, db):
        validate_sql("SELECT school_name FROM enrollment", conn=db)

    def test_accepts_aggregate_query(self, db):
        validate_sql(
            "SELECT school_name, SUM(student_count) FROM enrollment GROUP BY school_name",
            conn=db,
        )

    def test_accepts_join(self, db):
        validate_sql(
            "SELECT * FROM enrollment e JOIN attendance a ON e.school_name = a.school_name",
            conn=db,
        )


# ── Multi-statement attack prevention ──────────────────────────────────

class TestMultiStatement:
    """Multi-statement SQL injection should be blocked."""

    def test_semicolon_injection(self):
        """SELECT 1; DROP TABLE — forbidden token in second statement."""
        with pytest.raises(ValueError, match="drop"):
            # The whole string still contains 'drop' even if split by semicolon
            validate_sql("SELECT 1; DROP TABLE enrollment")

    def test_comment_then_drop(self):
        """SELECT 1 --\nDROP TABLE — forbidden token after comment."""
        with pytest.raises(ValueError, match="drop"):
            validate_sql("SELECT 1 -- harmless\nDROP TABLE enrollment")


# ── End-to-end: execute_safe ───────────────────────────────────────────

class TestExecuteSafe:
    """execute_safe() end-to-end pipeline."""

    def test_valid_query_returns_data(self, db):
        sql, df = execute_safe(db, "```sql\nSELECT COUNT(*) AS cnt FROM attendance\n```")
        assert sql == "SELECT COUNT(*) AS cnt FROM attendance"
        assert df.shape[0] == 1
        assert df["cnt"][0] > 0

    def test_invalid_sql_raises(self, db):
        with pytest.raises(ValueError):
            execute_safe(db, "```sql\nSELECT fake FROM nowhere\n```")

    def test_no_sql_blocks_uses_raw(self, db):
        sql, df = execute_safe(db, "SELECT 1 AS one")
        assert sql == "SELECT 1 AS one"
        assert df.shape == (1, 1)

    def test_limit_wrapping(self, db):
        """The LIMIT wrapper should cap results."""
        sql, df = execute_safe(db, "```sql\nSELECT * FROM attendance\n```")
        assert sql == "SELECT * FROM attendance"
        # Attendance has 11600 rows, but LIMIT 1000 caps it
        assert len(df) <= 1000