tests/test_safety.py

import pytest

from datapilot.safety import UnsafeCodeError, validate_generated_expression


def test_accepts_scalar_expression():
    tree = validate_generated_expression("(revenue - cost) / revenue")
    assert tree is not None


@pytest.mark.parametrize(
    "expression",
    [
        "__import__('os').system('whoami')",
        "open('secret.txt').read()",
        "(1).__class__.__mro__",
    ],
)
def test_blocks_unsafe_expression(expression):
    with pytest.raises(UnsafeCodeError):
        validate_generated_expression(expression)