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 = "

Hello& world

" raw += "" 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 = "

" + ("A" * 5000) + "

" 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"}