310 lines
11 KiB
Python
310 lines
11 KiB
Python
"""Tests for the formula parser (KTD7 security + parsing).
|
|
|
|
Test-first per U3 execution note: parser, security constraints, and cycle
|
|
detection tests are written before the engine/recalc worker.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from agentkit.bitable.formula.parser import (
|
|
MAX_AST_DEPTH,
|
|
FormulaDepthExceededError,
|
|
FormulaParseError,
|
|
FormulaSecurityError,
|
|
UnknownFunctionError,
|
|
evaluate_ast,
|
|
parse_formula,
|
|
)
|
|
|
|
ALLOWED = {"SUM", "AVG", "COUNT", "MIN", "MAX", "ABS", "ROUND", "IF", "LEN", "CONCAT"}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Parsing happy paths
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_parse_simple_arithmetic() -> None:
|
|
tree, mapping, _ = parse_formula("=1+2*3", ALLOWED)
|
|
assert mapping == {}
|
|
result = evaluate_ast(tree, {}, {})
|
|
assert result == 7
|
|
|
|
|
|
def test_parse_strips_equals_prefix() -> None:
|
|
tree1, _, _ = parse_formula("=1+1", ALLOWED)
|
|
tree2, _, _ = parse_formula("1+1", ALLOWED)
|
|
assert evaluate_ast(tree1, {}, {}) == evaluate_ast(tree2, {}, {}) == 2
|
|
|
|
|
|
def test_parse_field_reference() -> None:
|
|
tree, mapping, _ = parse_formula("={field_abc} + 1", ALLOWED)
|
|
assert "field_abc" in mapping.values()
|
|
# Safe name is prefixed with _f_
|
|
safe_name = next(k for k, v in mapping.items() if v == "field_abc")
|
|
result = evaluate_ast(tree, {safe_name: 41}, {})
|
|
assert result == 42
|
|
|
|
|
|
def test_parse_uuid_field_reference() -> None:
|
|
"""Field IDs are UUIDs with hyphens — must be substituted to safe names."""
|
|
fid = "a1b2c3d4-e5f6-7890-abcd-ef1234567890"
|
|
tree, mapping, _ = parse_formula(f"={{{fid}}} * 2", ALLOWED)
|
|
# The mapping should have a safe name → original UUID
|
|
assert fid in mapping.values()
|
|
# Evaluate using the safe name (prefixed with _f_)
|
|
safe_name = next(k for k, v in mapping.items() if v == fid)
|
|
assert safe_name.startswith("_f_")
|
|
result = evaluate_ast(tree, {safe_name: 21}, {})
|
|
assert result == 42
|
|
|
|
|
|
def test_parse_string_concatenation() -> None:
|
|
tree, _, _ = parse_formula('="hello" + " " + "world"', ALLOWED)
|
|
assert evaluate_ast(tree, {}, {}) == "hello world"
|
|
|
|
|
|
def test_parse_conditional_ifexp() -> None:
|
|
tree, _, _ = parse_formula("=1 if True else 2", ALLOWED)
|
|
assert evaluate_ast(tree, {}, {}) == 1
|
|
|
|
|
|
def test_parse_comparison() -> None:
|
|
tree, mapping, _ = parse_formula("={f} > 5", ALLOWED)
|
|
safe_name = next(k for k, v in mapping.items() if v == "f")
|
|
assert evaluate_ast(tree, {safe_name: 10}, {}) is True
|
|
assert evaluate_ast(tree, {safe_name: 3}, {}) is False
|
|
|
|
|
|
def test_parse_boolean_ops() -> None:
|
|
tree, _, _ = parse_formula("=True and False", ALLOWED)
|
|
assert evaluate_ast(tree, {}, {}) is False
|
|
tree2, _, _ = parse_formula("=True or False", ALLOWED)
|
|
assert evaluate_ast(tree2, {}, {}) is True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Function calls
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_parse_function_call_sum() -> None:
|
|
tree, mapping, _ = parse_formula("=SUM({f1})", ALLOWED)
|
|
safe_name = next(k for k, v in mapping.items() if v == "f1")
|
|
result = evaluate_ast(tree, {safe_name: [1, 2, 3]}, {"SUM": sum})
|
|
assert result == 6
|
|
|
|
|
|
def test_parse_function_call_concat() -> None:
|
|
tree, mapping, _ = parse_formula('=CONCAT({f1}, "-", {f2})', ALLOWED)
|
|
safe_f1 = next(k for k, v in mapping.items() if v == "f1")
|
|
safe_f2 = next(k for k, v in mapping.items() if v == "f2")
|
|
result = evaluate_ast(
|
|
tree, {safe_f1: "a", safe_f2: "b"}, {"CONCAT": lambda *a: "".join(str(x) for x in a)}
|
|
)
|
|
assert result == "a-b"
|
|
|
|
|
|
def test_parse_nested_function_calls() -> None:
|
|
tree, _, _ = parse_formula("=ABS(-5) + ROUND(3.7, 0)", ALLOWED)
|
|
funcs = {"ABS": abs, "ROUND": round}
|
|
result = evaluate_ast(tree, {}, funcs)
|
|
assert result == 9 # 5 + 4
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# KTD7 Security — disallowed nodes
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_security_rejects_attribute_access() -> None:
|
|
"""__import__('os') is rejected — it's a Call to an unregistered function.
|
|
(Attribute access like os.system would be caught by the Attribute node check,
|
|
but __import__ is caught earlier as an unknown function.)"""
|
|
with pytest.raises((FormulaSecurityError, UnknownFunctionError)):
|
|
parse_formula("=__import__('os')", ALLOWED)
|
|
|
|
|
|
def test_security_rejects_attribute_chain() -> None:
|
|
"""Attribute access like ''.join([]) is rejected by the Attribute node check."""
|
|
with pytest.raises(FormulaSecurityError):
|
|
parse_formula("=''.join([])", ALLOWED)
|
|
|
|
|
|
def test_security_rejects_lambda() -> None:
|
|
with pytest.raises(FormulaSecurityError):
|
|
parse_formula("=(lambda: 1)()", ALLOWED)
|
|
|
|
|
|
def test_security_rejects_subscript() -> None:
|
|
with pytest.raises(FormulaSecurityError):
|
|
parse_formula("=[1,2,3][0]", ALLOWED)
|
|
|
|
|
|
def test_security_rejects_assignment() -> None:
|
|
"""Assignment is a statement, not an expression — rejected at parse stage."""
|
|
with pytest.raises((FormulaSecurityError, FormulaParseError)):
|
|
parse_formula("=x = 1", ALLOWED)
|
|
|
|
|
|
def test_unknown_function_rejected() -> None:
|
|
with pytest.raises(UnknownFunctionError):
|
|
parse_formula("=UNKNOWN(1)", ALLOWED)
|
|
|
|
|
|
def test_eval_function_rejected_if_not_registered() -> None:
|
|
"""eval is not in the registry → UnknownFunctionError."""
|
|
with pytest.raises(UnknownFunctionError):
|
|
parse_formula("=eval('1+1')", ALLOWED)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Error paths
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_parse_error_unbalanced_parens() -> None:
|
|
with pytest.raises(FormulaParseError):
|
|
parse_formula("=(1+2", ALLOWED)
|
|
|
|
|
|
def test_parse_error_empty_formula() -> None:
|
|
with pytest.raises(FormulaParseError):
|
|
parse_formula("=", ALLOWED)
|
|
|
|
|
|
def test_parse_error_empty_string() -> None:
|
|
with pytest.raises(FormulaParseError):
|
|
parse_formula("", ALLOWED)
|
|
|
|
|
|
def test_evaluate_unknown_field_ref_raises() -> None:
|
|
tree, _, _ = parse_formula("={nonexistent} + 1", ALLOWED)
|
|
with pytest.raises(FormulaParseError, match="Unknown field reference"):
|
|
evaluate_ast(tree, {}, {})
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Mixed aggregate + row context
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_mixed_aggregate_and_row_context() -> None:
|
|
"""={f1} + SUM({f2}) — row f1 + column f2 sum."""
|
|
tree, mapping, _ = parse_formula("={f1} + SUM({f2})", ALLOWED)
|
|
safe_f1 = next(k for k, v in mapping.items() if v == "f1")
|
|
safe_f2 = next(k for k, v in mapping.items() if v == "f2")
|
|
# f1 is a row value (scalar), f2 is a column value (list)
|
|
result = evaluate_ast(tree, {safe_f1: 10, safe_f2: [1, 2, 3]}, {"SUM": sum})
|
|
assert result == 16 # 10 + 6
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# KTD8: AST depth limit
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_depth_limit_allows_normal_formula() -> None:
|
|
"""Normal formulas should pass the depth check."""
|
|
tree, _, _ = parse_formula("=1+2+3+4+5", ALLOWED)
|
|
assert tree is not None
|
|
|
|
|
|
def test_depth_limit_rejects_deeply_nested() -> None:
|
|
"""Deeply nested expressions should be rejected (KTD8)."""
|
|
# Build a formula deeper than MAX_AST_DEPTH: nested additions
|
|
# Each "x + " adds 1 level of BinOp depth
|
|
deep_formula = "=" + " + ".join(["1"] * (MAX_AST_DEPTH + 5))
|
|
with pytest.raises(FormulaDepthExceededError, match="depth"):
|
|
parse_formula(deep_formula, ALLOWED)
|
|
|
|
|
|
def test_depth_limit_allows_at_boundary() -> None:
|
|
"""Formula at exactly the limit should pass."""
|
|
# Build a formula with depth close to but not exceeding MAX_AST_DEPTH
|
|
# A simple sum of N numbers has depth ~N, but ast groups them differently.
|
|
# A safer test: shallow formula always passes.
|
|
tree, _, _ = parse_formula("=SUM({f1}) + IF({f2}, {f3}, {f4})", ALLOWED)
|
|
assert tree is not None
|
|
|
|
|
|
def test_list_literal_evaluated() -> None:
|
|
"""List literals [1, 2, 3] should be evaluated as Python lists."""
|
|
tree, _, _ = parse_formula("=[1, 2, 3]", ALLOWED)
|
|
result = evaluate_ast(tree, {}, {})
|
|
assert result == [1, 2, 3]
|
|
|
|
|
|
def test_tuple_literal_evaluated() -> None:
|
|
"""Tuple literals (1, 2, 3) should be evaluated as Python tuples."""
|
|
tree, _, _ = parse_formula("=(1, 2, 3)", ALLOWED)
|
|
result = evaluate_ast(tree, {}, {})
|
|
assert result == (1, 2, 3)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# U3: Cross-table references {rel.target}
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_parse_cross_table_reference() -> None:
|
|
"""{rel_field.target_field} is parsed into a single Name node (KTD-2)."""
|
|
tree, field_mapping, cross_table_mapping = parse_formula("={f_rel.f_target}", ALLOWED)
|
|
# Single-field mapping should be empty
|
|
assert field_mapping == {}
|
|
# Cross-table mapping should have one entry
|
|
assert len(cross_table_mapping) == 1
|
|
safe_name, (rel_id, target_id) = next(iter(cross_table_mapping.items()))
|
|
assert rel_id == "f_rel"
|
|
assert target_id == "f_target"
|
|
assert safe_name.startswith("_ct_")
|
|
assert "__" in safe_name # separator between rel and target
|
|
|
|
|
|
def test_parse_cross_table_with_aggregate() -> None:
|
|
"""=SUM({rel.target}) — cross-table ref inside an aggregate function."""
|
|
tree, _, cross_table_mapping = parse_formula("=SUM({f_rel.f_target})", ALLOWED)
|
|
assert len(cross_table_mapping) == 1
|
|
safe_name, (rel_id, target_id) = next(iter(cross_table_mapping.items()))
|
|
assert rel_id == "f_rel"
|
|
assert target_id == "f_target"
|
|
# Should evaluate correctly when given a list of values
|
|
result = evaluate_ast(tree, {safe_name: [1, 2, 3]}, {"SUM": sum})
|
|
assert result == 6
|
|
|
|
|
|
def test_parse_mixed_single_and_cross_table_refs() -> None:
|
|
"""={f_local} + SUM({f_rel.f_target}) — mixed row + cross-table aggregate."""
|
|
tree, field_mapping, cross_table_mapping = parse_formula(
|
|
"={f_local} + SUM({f_rel.f_target})", ALLOWED
|
|
)
|
|
assert "f_local" in field_mapping.values()
|
|
assert len(cross_table_mapping) == 1
|
|
safe_local = next(k for k, v in field_mapping.items() if v == "f_local")
|
|
safe_cross = next(iter(cross_table_mapping.keys()))
|
|
result = evaluate_ast(
|
|
tree, {safe_local: 10, safe_cross: [1, 2, 3]}, {"SUM": sum}
|
|
)
|
|
assert result == 16
|
|
|
|
|
|
def test_parse_cross_table_uuid_field_ids() -> None:
|
|
"""Cross-table refs work with UUID-style field IDs containing hyphens."""
|
|
rel_id = "a1b2c3d4-e5f6-7890-abcd-ef1234567890"
|
|
target_id = "11111111-2222-3333-4444-555555555555"
|
|
tree, _, cross_table_mapping = parse_formula(f"={{{rel_id}.{target_id}}}", ALLOWED)
|
|
assert len(cross_table_mapping) == 1
|
|
_, (parsed_rel, parsed_target) = next(iter(cross_table_mapping.items()))
|
|
assert parsed_rel == rel_id
|
|
assert parsed_target == target_id
|
|
|
|
|
|
def test_cross_table_ref_does_not_match_single_field() -> None:
|
|
"""{rel.target} must not be partially matched by _FIELD_REF_RE."""
|
|
# The rel part should not leak into field_mapping
|
|
tree, field_mapping, cross_table_mapping = parse_formula("={f_rel.f_target}", ALLOWED)
|
|
assert field_mapping == {}
|
|
assert len(cross_table_mapping) == 1
|