280 lines
11 KiB
Python
280 lines
11 KiB
Python
"""Paid-freeze tests: default auto is free-only, pin still works,
|
|
unfreeze is explicit (WEBSEARCH_ENABLE_PAID=1, exact value "1").
|
|
|
|
No network, no real keys: every HTTP call goes through a FakeHttp,
|
|
and the fake key values are deliberately distinctive so a leak into
|
|
an error string or a tool result would be caught by assertion.
|
|
"""
|
|
|
|
import json
|
|
import socket
|
|
import sys
|
|
import unittest
|
|
import urllib.error
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
try:
|
|
from cloud.websearch import (
|
|
DEFAULT_CHAIN, DEFAULT_FREE_CHAIN, HttpClient, WebSearchError,
|
|
available_providers, paid_search_enabled,
|
|
resolve_provider_chain, search)
|
|
from cloud.websearch_mcp import handle_request
|
|
except ImportError:
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
|
from cloud.websearch import (
|
|
DEFAULT_CHAIN, DEFAULT_FREE_CHAIN, HttpClient, WebSearchError,
|
|
available_providers, paid_search_enabled,
|
|
resolve_provider_chain, search)
|
|
from cloud.websearch_mcp import handle_request
|
|
|
|
_FAKE_KEYS = {
|
|
"PERPLEXITY_API_KEY": "pplx-SECRET-111",
|
|
"EXA_API_KEY": "exa-SECRET-222",
|
|
"TAVILY_API_KEY": "tvly-SECRET-333",
|
|
"BRAVE_API_KEY": "brave-SECRET-444",
|
|
}
|
|
|
|
_DDG_HTML = (
|
|
'<a class="result__a" href="//duckduckgo.com/l/?uddg='
|
|
'https%3A%2F%2Fexample.com%2Fddg">DDG Title</a>'
|
|
'<a class="result__snippet">ddg snippet</a>')
|
|
|
|
|
|
class FakeHttp:
|
|
"""post_json/get keyed by URL substring; failures raise like HTTP."""
|
|
|
|
def __init__(self, routes):
|
|
self.routes = routes
|
|
self.seen = []
|
|
|
|
def post_json(self, url, headers, payload):
|
|
self.seen.append(("POST", url, payload))
|
|
return self._route(url)
|
|
|
|
def get(self, url, headers=None):
|
|
self.seen.append(("GET", url, headers))
|
|
return self._route(url)
|
|
|
|
def _route(self, url):
|
|
for needle, body, status in self.routes:
|
|
if needle in url:
|
|
if status != 200:
|
|
raise WebSearchError("HTTP %s" % status)
|
|
if isinstance(body, dict):
|
|
return 200, json.dumps(body).encode(), "application/json"
|
|
return 200, body.encode(), "text/html"
|
|
raise WebSearchError("HTTP 404")
|
|
|
|
|
|
def _searxng_body():
|
|
return {"results": [
|
|
{"title": "SearXNG Hit", "url": "https://example.com/sx",
|
|
"content": "searxng snippet"}]}
|
|
|
|
|
|
def _paid_env(**extra):
|
|
env = dict(_FAKE_KEYS)
|
|
env["SEARXNG_URL"] = "http://searx.test"
|
|
env.update(extra)
|
|
return env
|
|
|
|
|
|
def _paid_hosts_seen(http):
|
|
hosts = []
|
|
for _method, url, _payload in http.seen:
|
|
for paid in ("perplexity", "exa.ai", "tavily",
|
|
"search.brave.com"):
|
|
if paid in url:
|
|
hosts.append(paid)
|
|
return hosts
|
|
|
|
|
|
class FrozenDefaultTestCase(unittest.TestCase):
|
|
def test_default_auto_excludes_paid_even_with_all_keys(self):
|
|
env = _paid_env()
|
|
self.assertEqual(available_providers(env),
|
|
["searxng", "duckduckgo"])
|
|
self.assertEqual(resolve_provider_chain("auto", env),
|
|
["searxng", "duckduckgo"])
|
|
http = FakeHttp([("searx", _searxng_body(), 200)])
|
|
out = search("q", provider="auto", env=env, http=http)
|
|
self.assertEqual(out["provider"], "searxng")
|
|
self.assertEqual(_paid_hosts_seen(http), [])
|
|
|
|
def test_default_auto_without_searxng_url_is_duckduckgo_only(self):
|
|
env = dict(_FAKE_KEYS)
|
|
self.assertEqual(resolve_provider_chain("auto", env),
|
|
["duckduckgo"])
|
|
http = FakeHttp([("duckduckgo", _DDG_HTML, 200)])
|
|
out = search("q", provider="auto", env=env, http=http)
|
|
self.assertEqual(out["provider"], "duckduckgo")
|
|
self.assertEqual(_paid_hosts_seen(http), [])
|
|
|
|
def test_searxng_failure_never_falls_back_to_paid(self):
|
|
http = FakeHttp([("searx", {}, 500),
|
|
("duckduckgo", "<html>challenge</html>", 200)])
|
|
with self.assertRaises(WebSearchError) as ctx:
|
|
search("q", provider="auto", env=_paid_env(), http=http)
|
|
self.assertEqual(
|
|
[a["provider"] for a in ctx.exception.attempted],
|
|
["searxng", "duckduckgo"])
|
|
self.assertEqual(_paid_hosts_seen(http), [])
|
|
# No fabricated result: the error names both free stages.
|
|
self.assertIn("searxng", str(ctx.exception))
|
|
self.assertIn("duckduckgo", str(ctx.exception))
|
|
|
|
def test_duckduckgo_fallback_failure_reports_clearly(self):
|
|
# DuckDuckGo under the VM/host egress returns a 202-style
|
|
# robot-challenge page: the parser yields no results, which
|
|
# must surface as an honest "no results" failure for that
|
|
# stage, never as a fabricated hit.
|
|
http = FakeHttp([("searx", {}, 500),
|
|
("duckduckgo", "<html>anomaly</html>", 200)])
|
|
with self.assertRaises(WebSearchError) as ctx:
|
|
search("q", provider="auto", env=_paid_env(), http=http)
|
|
attempted = {a["provider"]: a for a in ctx.exception.attempted}
|
|
self.assertEqual(attempted["duckduckgo"]["error"], "no results")
|
|
self.assertFalse(attempted["duckduckgo"]["ok"])
|
|
for secret in _FAKE_KEYS.values():
|
|
self.assertNotIn(secret, str(ctx.exception))
|
|
|
|
|
|
class ExplicitPinTestCase(unittest.TestCase):
|
|
def test_explicit_pin_paid_still_works_while_frozen(self):
|
|
http = FakeHttp([
|
|
("api.tavily.com",
|
|
{"results": [{"title": "T", "url": "https://example.com/t",
|
|
"content": "tavily content"}]}, 200),
|
|
])
|
|
out = search("q", provider="tavily", env=_paid_env(), http=http)
|
|
self.assertEqual(out["provider"], "tavily")
|
|
self.assertEqual(out["results"][0]["snippet"], "tavily content")
|
|
|
|
def test_explicit_pin_paid_failure_does_not_fall_back(self):
|
|
http = FakeHttp([("perplexity", {}, 500),
|
|
("api.exa.ai", {"results": [
|
|
{"title": "E", "url": "https://example.com/e",
|
|
"text": "x"}]}, 200)])
|
|
with self.assertRaises(WebSearchError):
|
|
search("q", provider="perplexity", env=_paid_env(), http=http)
|
|
self.assertEqual(_paid_hosts_seen(http), ["perplexity"])
|
|
|
|
def test_frozen_default_applies_to_scope_code_too(self):
|
|
body = {"results": [
|
|
{"title": "GH", "url": "https://github.com/a/b",
|
|
"content": "code hit"},
|
|
{"title": "Other", "url": "https://example.com/other",
|
|
"content": "off-scope"}]}
|
|
http = FakeHttp([("searx", body, 200)])
|
|
out = search("q", provider="auto", env=_paid_env(),
|
|
http=http, scope="code")
|
|
self.assertEqual(out["provider"], "searxng")
|
|
self.assertEqual(
|
|
[r["url"] for r in out["results"]], ["https://github.com/a/b"])
|
|
self.assertEqual(_paid_hosts_seen(http), [])
|
|
|
|
|
|
class UnfreezeSwitchTestCase(unittest.TestCase):
|
|
def test_switch_off_by_default_and_non_1_values_do_not_unfreeze(self):
|
|
self.assertFalse(paid_search_enabled({}))
|
|
for value in ("0", "true", "yes", "", "on"):
|
|
env = _paid_env(WEBSEARCH_ENABLE_PAID=value)
|
|
self.assertFalse(paid_search_enabled(env), value)
|
|
self.assertEqual(resolve_provider_chain("auto", env),
|
|
["searxng", "duckduckgo"])
|
|
|
|
def test_switch_on_restores_legacy_chain_and_fallback(self):
|
|
env = _paid_env(WEBSEARCH_ENABLE_PAID="1")
|
|
self.assertTrue(paid_search_enabled(env))
|
|
self.assertEqual(resolve_provider_chain("auto", env),
|
|
list(DEFAULT_CHAIN))
|
|
# The pre-freeze perplexity -> exa fallback, preserved
|
|
# verbatim under the explicit switch (plus the no-leak
|
|
# assertion from the old default-chain test).
|
|
http = FakeHttp([
|
|
("perplexity", {}, 500),
|
|
("api.exa.ai", {"results": [
|
|
{"title": "Exa Hit", "url": "https://example.com/exa",
|
|
"text": "exa snippet", "score": 0.9}]}, 200),
|
|
])
|
|
out = search("test query", provider="auto", env=env, http=http)
|
|
self.assertEqual(out["provider"], "exa")
|
|
self.assertEqual(out["results"][0]["url"],
|
|
"https://example.com/exa")
|
|
self.assertEqual(out["attempted"][0]["provider"], "perplexity")
|
|
self.assertFalse(out["attempted"][0]["ok"])
|
|
blob = json.dumps(out)
|
|
for secret in _FAKE_KEYS.values():
|
|
self.assertNotIn(secret, blob)
|
|
|
|
def test_constants_state_the_contract(self):
|
|
self.assertEqual(DEFAULT_FREE_CHAIN, ("searxng", "duckduckgo"))
|
|
self.assertEqual(
|
|
DEFAULT_CHAIN,
|
|
("perplexity", "exa", "tavily", "brave", "searxng",
|
|
"duckduckgo"))
|
|
|
|
|
|
class TimeoutClarityTestCase(unittest.TestCase):
|
|
def test_http_client_maps_timeout_to_explicit_error(self):
|
|
client = HttpClient(timeout=20)
|
|
with mock.patch("urllib.request.urlopen",
|
|
side_effect=socket.timeout("timed out")):
|
|
with self.assertRaises(WebSearchError) as ctx:
|
|
client.get("http://searx.test/search?q=x")
|
|
self.assertIn("timeout after 20s", str(ctx.exception))
|
|
|
|
with mock.patch(
|
|
"urllib.request.urlopen",
|
|
side_effect=urllib.error.URLError(
|
|
socket.timeout("timed out"))):
|
|
with self.assertRaises(WebSearchError) as ctx:
|
|
client.get("http://searx.test/search?q=x")
|
|
self.assertIn("timeout after 20s", str(ctx.exception))
|
|
|
|
def test_http_client_does_not_relabel_dns_failure_as_timeout(self):
|
|
client = HttpClient(timeout=20)
|
|
with mock.patch(
|
|
"urllib.request.urlopen",
|
|
side_effect=urllib.error.URLError(
|
|
"Temporary failure in name resolution")):
|
|
with self.assertRaises(WebSearchError) as ctx:
|
|
client.get("http://searx.test/search?q=x")
|
|
self.assertNotIn("timeout", str(ctx.exception))
|
|
self.assertIn("name resolution", str(ctx.exception))
|
|
|
|
def test_search_reports_timeout_as_error_not_no_results(self):
|
|
class TimeoutHttp(FakeHttp):
|
|
def get(self, url, headers=None):
|
|
self.seen.append(("GET", url, headers))
|
|
if "searx" in url:
|
|
raise WebSearchError(
|
|
"timeout after 20s waiting for HTTP response")
|
|
return self._route(url)
|
|
|
|
http = TimeoutHttp([("duckduckgo", _DDG_HTML, 200)])
|
|
out = search("q", provider="auto", env=_paid_env(), http=http)
|
|
self.assertEqual(out["provider"], "duckduckgo")
|
|
attempted = {a["provider"]: a for a in out["attempted"]}
|
|
self.assertIn("timeout after 20s",
|
|
attempted["searxng"]["error"])
|
|
self.assertNotEqual(attempted["searxng"]["error"], "no results")
|
|
|
|
|
|
class McpDescriptionTestCase(unittest.TestCase):
|
|
def test_tools_list_describes_frozen_default_and_unfreeze(self):
|
|
resp = handle_request(
|
|
{"jsonrpc": "2.0", "id": 1, "method": "tools/list"})
|
|
tool = [t for t in resp["result"]["tools"]
|
|
if t["name"] == "web_search"][0]
|
|
desc = tool["description"]
|
|
self.assertIn("SearXNG", desc)
|
|
self.assertIn("DuckDuckGo", desc)
|
|
self.assertIn("WEBSEARCH_ENABLE_PAID=1", desc)
|
|
self.assertIn("paid", desc.lower())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|