Spaces:
Sleeping
Sleeping
| 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& <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"} | |