lambdagent/tests/test_extensions.py

"""Tests for lambdagent extensions: Par, Route, Memory, Guard."""
import pytest
from lambdagent.core import Context, RouteError, ValidationError
from lambdagent.primitives import Tool
from lambdagent.extensions import Par, Route, Memory, Guard


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


def test_par():
    a = t("x2", lambda x: int(x) * 2)
    b = t("x3", lambda x: int(x) * 3)
    p = Par(a, b)
    assert p("4") == (8, 12)


def test_par_operator():
    a = t("a", lambda x: "A")
    b = t("b", lambda x: "B")
    c = t("c", lambda x: "C")
    p = a | b | c
    assert p("x") == ("A", "B", "C")


def test_route():
    cls = t("cls", lambda x: "math" if "+" in x else "text")
    math_agent = t("math", lambda x: "math_result")
    text_agent = t("text", lambda x: "text_result")
    router = Route(cls, {"math": math_agent, "text": text_agent})
    assert router("1+1") == "math_result"
    assert router("hello") == "text_result"


def test_route_default():
    cls = t("cls", lambda x: "unknown")
    default = t("default", lambda x: "default_result")
    router = Route(cls, {"a": t("a", lambda x: "A")}, default=default)
    assert router("x") == "default_result"


def test_route_no_match():
    cls = t("cls", lambda x: "missing")
    router = Route(cls, {"a": t("a", lambda x: "A")})
    with pytest.raises(RouteError):
        router("x")


def test_memory():
    agent = t("echo", lambda x: x)
    mem = Memory(agent, store={"user": "Alice"})
    result = mem("hello")
    assert "Alice" in result
    assert "hello" in result


def test_memory_remember_forget():
    agent = t("echo", lambda x: x)
    mem = Memory(agent)
    mem.remember("k", "v")
    assert mem.store["k"] == "v"
    mem.forget("k")
    assert "k" not in mem.store


def test_guard_pass():
    agent = t("gen", lambda x: "long enough output text here")
    g = Guard(agent, validator=lambda x: len(x) > 5)
    assert g("x") == "long enough output text here"


def test_guard_fail():
    agent = t("gen", lambda x: "hi")
    g = Guard(agent, validator=lambda x: len(x) > 100, retry=1)
    with pytest.raises(ValidationError):
        g("x")


def test_guard_on_fail():
    agent = t("gen", lambda x: "short")
    g = Guard(agent, validator=lambda x: len(x) > 100, retry=0, on_fail=lambda x: "fallback")
    assert g("x") == "fallback"