observable-agent / tests /test_tools.py
sukhrobnurali's picture
Deploy Observable Agent Space
f019486 verified
Raw
History Blame Contribute Delete
6.17 kB
from agent import tools
def test_calculator_arithmetic():
assert tools.calculator("2 + 3 * 4") == "14"
assert tools.calculator("(1 + 2) ** 3") == "27"
assert tools.calculator("17 % 5") == "2"
def test_calculator_rejects_imports_and_builtins():
assert tools.calculator("__import__('os').getcwd()").startswith("ERROR")
assert tools.calculator("eval('1+1')").startswith("ERROR")
def test_calculator_strips_dangerous_symbols():
# open/eval/etc. are removed from the evaluator symtable, so even a bare reference is undefined.
out = tools.calculator("open")
assert out.startswith("ERROR") and "not defined" in out
def test_calculator_rejects_non_numeric_input():
# Strings/lists/tuples/dicts are rejected outright: blocks quoted file-path access and the
# unbounded sequence-multiplication memory DoS that asteval does not guard.
assert tools.calculator("open('requirements.txt')").startswith("ERROR")
assert tools.calculator("[0]*10**9").startswith("ERROR")
assert tools.calculator("(0,)*10**9").startswith("ERROR")
assert tools.calculator("'a'*10**9").startswith("ERROR")
assert tools.calculator("{0:0}").startswith("ERROR")
def test_calculator_degrades_on_huge_results():
# Unbounded integer growth (huge exponent string conversion, removed factorial) must return
# ERROR, not crash the process or hang.
assert tools.calculator("10**10000").startswith("ERROR")
assert tools.calculator("factorial(100000)").startswith("ERROR")
def test_calculator_rejects_too_long():
assert tools.calculator("1+" * 500 + "1").startswith("ERROR")
def test_calculator_division_by_zero_is_error():
assert tools.calculator("1/0").startswith("ERROR")
def test_strip_html_removes_tags_and_scripts():
raw = "<html><head><style>.x{}</style></head><body><p>Hello&amp; <b>world</b></p>"
raw += "<script>alert(1)</script></body></html>"
cleaned = tools.strip_html(raw)
assert "Hello& world" in cleaned
assert "alert" not in cleaned
assert "<" not in cleaned
def test_web_get_rejects_non_http_scheme():
assert tools.web_get("file:///etc/passwd").startswith("ERROR")
assert tools.web_get("ftp://example.com").startswith("ERROR")
def test_web_get_caps_length_and_uses_injected_client():
class FakeResp:
text = "<p>" + ("A" * 5000) + "</p>"
def raise_for_status(self):
return None
class FakeClient:
def __init__(self, *a, **k):
pass
def __enter__(self):
return self
def __exit__(self, *a):
return False
def get(self, url):
return FakeResp()
out = tools.web_get("https://example.com", client_factory=FakeClient, host_check=lambda h: True)
assert len(out) <= tools.config.WEB_MAX_CHARS
assert out.startswith("A")
def test_web_get_handles_fetch_error():
class BoomClient:
def __init__(self, *a, **k):
pass
def __enter__(self):
return self
def __exit__(self, *a):
return False
def get(self, url):
raise ConnectionError("boom")
assert tools.web_get("https://example.com", client_factory=BoomClient, host_check=lambda h: True).startswith("ERROR")
def _fake_wiki(summary, exists=True):
class FakePage:
def __init__(self):
self.summary = summary
def exists(self):
return exists
class FakeWiki:
def page(self, query):
return FakePage()
return FakeWiki()
def test_wikipedia_search_returns_capped_summary():
long_summary = "Paris is the capital of France. " * 50
out = tools.wikipedia_search("Paris", wiki=_fake_wiki(long_summary))
assert out.startswith("Paris is the capital")
assert len(out) <= tools.config.WIKI_SUMMARY_CHARS
def test_wikipedia_search_missing_page_returns_no_results():
out = tools.wikipedia_search("Asdkjfh", wiki=_fake_wiki("", exists=False))
assert out.startswith("NO_RESULTS")
def test_wikipedia_search_handles_error():
class BoomWiki:
def page(self, query):
raise RuntimeError("network down")
assert tools.wikipedia_search("x", wiki=BoomWiki()).startswith("ERROR")
def test_host_is_public_blocks_private_and_loopback():
def private(host, *a, **k):
return [(2, 1, 6, "", ("10.0.0.5", 0))]
def public(host, *a, **k):
return [(2, 1, 6, "", ("93.184.216.34", 0))]
assert tools._host_is_public("internal", resolver=private) is False
assert tools._host_is_public("example.com", resolver=public) is True
assert tools._host_is_public("", resolver=public) is False
def test_web_get_blocks_non_public_host():
out = tools.web_get("https://metadata.internal", host_check=lambda h: False)
assert out.startswith("ERROR") and "host" in out.lower()
def test_web_get_revalidates_redirect_target():
class RedirResp:
is_redirect = True
headers = {"location": "http://169.254.169.254/"}
def raise_for_status(self):
return None
class RedirClient:
def __init__(self, *a, **k):
pass
def __enter__(self):
return self
def __exit__(self, *a):
return False
def get(self, url):
return RedirResp()
out = tools.web_get("https://safe.example.com", client_factory=RedirClient,
host_check=lambda h: h != "169.254.169.254")
assert out.startswith("ERROR")
def test_registry_has_three_tools_and_callables():
callables = tools.tool_callables()
assert set(callables) == {"calculator", "wikipedia_search", "web_get"}
assert callables["calculator"]("2+2") == "4"
def test_declarations_is_single_tool_with_three_functions():
decls = tools.declarations()
assert len(decls) == 1
names = {fd.name for fd in decls[0].function_declarations}
assert names == {"calculator", "wikipedia_search", "web_get"}