Spaces:
Sleeping
Sleeping
File size: 6,171 Bytes
f019486 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 | 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"}
|