lambdagent/tests/test_primitives.py

"""Tests for lambdagent primitives: Compose, If, Loop, Pair, Fst, Snd, Tool."""
import pytest
from lambdagent.core import Context
from lambdagent.primitives import Compose, If, Loop, Pair, Fst, Snd, Tool


def make_tool(name, fn):
    return Tool(name, fn)


def test_tool_apply():
    t = make_tool("double", lambda x: int(x) * 2)
    assert t("3") == 6


def test_tool_trace():
    ctx = Context()
    t = make_tool("inc", lambda x: int(x) + 1)
    t.apply("5", ctx)
    assert len(ctx.trace) == 1
    assert ctx.trace[0].output == 6


def test_compose():
    a = make_tool("add1", lambda x: int(x) + 1)
    b = make_tool("mul2", lambda x: int(x) * 2)
    pipe = Compose(a, b)
    assert pipe("3") == 8  # (3+1)*2


def test_compose_operator():
    a = make_tool("a", lambda x: x + "A")
    b = make_tool("b", lambda x: x + "B")
    c = make_tool("c", lambda x: x + "C")
    pipe = a >> b >> c
    assert pipe("") == "ABC"


def test_if_truthy():
    t = make_tool("yes", lambda x: "YES")
    f = make_tool("no", lambda x: "NO")
    cond = If(lambda x: x == "go", t, f)
    assert cond("go") == "YES"
    assert cond("stop") == "NO"


def test_if_term_cond():
    cond_term = make_tool("check", lambda x: "TRUE" if len(x) > 3 else "FALSE")
    t = make_tool("long", lambda x: "long")
    f = make_tool("short", lambda x: "short")
    branch = If(cond_term, t, f)
    assert branch("hello") == "long"
    assert branch("hi") == "short"


def test_loop():
    counter = {"n": 0}
    def step(x):
        counter["n"] += 1
        return int(x) + 1
    body = make_tool("inc", step)
    loop = Loop(body, condition=lambda r, s: int(r) >= 5, max_steps=20)
    result = loop("0")
    assert result == 5
    assert counter["n"] == 5


def test_loop_max_steps():
    body = make_tool("noop", lambda x: x)
    loop = Loop(body, condition=lambda r, s: False, max_steps=3)
    loop("x")  # Should not infinite loop


def test_pair():
    a = make_tool("upper", lambda x: x.upper())
    b = make_tool("lower", lambda x: x.lower())
    p = Pair(a, b)
    assert p("Hello") == ("HELLO", "hello")


def test_fst_snd():
    assert Fst()(("a", "b")) == "a"
    assert Snd()(("a", "b")) == "b"