diff --git a/tests/test_search_router.py b/tests/test_search_router.py new file mode 100644 index 0000000..c201f57 --- /dev/null +++ b/tests/test_search_router.py @@ -0,0 +1,125 @@ +import pytest +from fastapi.testclient import TestClient +from fastapi import FastAPI +from axon.models import SearchResult + + +def _mock_result(url="https://ex.com", title="T", content="C", score=0.8): + return SearchResult(title=title, url=url, content=content, score=score) + + +async def _async(val): + return val + + +@pytest.fixture +def client(monkeypatch): + from axon.services.cache import InMemoryCache + cache = InMemoryCache() + + async def async_cache(): + return cache + + monkeypatch.setattr("axon.routers.search.get_cache", async_cache) + + results_5 = [_mock_result(f"https://ex{i}.com", f"T{i}", f"C{i}") for i in range(5)] + + async def fake_brave(*a, **kw): + return results_5 + + async def fake_ddg(*a, **kw): + return [] + + async def fake_rerank(q, c): + return [(url, text, 0.8) for url, text in c] + + monkeypatch.setattr("axon.routers.search.search_brave", fake_brave) + monkeypatch.setattr("axon.routers.search.search_ddg", fake_ddg) + monkeypatch.setattr("axon.routers.search.rerank_results", fake_rerank) + + from axon.routers.search import router + app = FastAPI() + app.include_router(router) + return TestClient(app) + + +def test_search_basic_returns_200(client): + resp = client.post("/search", json={"query": "test"}) + assert resp.status_code == 200 + data = resp.json() + assert "results" in data + assert "response_time" in data + assert data["query"] == "test" + + +def test_search_deduplicates_urls(): + from axon.routers.search import deduplicate_results + r1 = _mock_result("https://x.com") + r2 = _mock_result("https://x.com") + r3 = _mock_result("https://y.com") + deduped = deduplicate_results([r1, r2, r3]) + assert len(deduped) == 2 + urls = [r.url for r in deduped] + assert urls.count("https://x.com") == 1 + + +def test_search_503_if_all_sources_fail(monkeypatch): + from axon.services.cache import InMemoryCache + cache = InMemoryCache() + + async def async_cache(): + return cache + + monkeypatch.setattr("axon.routers.search.get_cache", async_cache) + + async def empty(*a, **kw): + return [] + + monkeypatch.setattr("axon.routers.search.search_brave", empty) + monkeypatch.setattr("axon.routers.search.search_ddg", empty) + monkeypatch.setattr("axon.routers.search.search_serper", empty) + + from axon.routers.search import router + app = FastAPI() + app.include_router(router) + c = TestClient(app, raise_server_exceptions=False) + resp = c.post("/search", json={"query": "test"}) + assert resp.status_code == 503 + + +def test_search_whitelist_filter(monkeypatch): + from axon.services.cache import InMemoryCache + cache = InMemoryCache() + + async def async_cache(): + return cache + + monkeypatch.setattr("axon.routers.search.get_cache", async_cache) + + results = [ + _mock_result("https://allowed.com/page"), + _mock_result("https://blocked.com/page"), + ] + + async def fake_brave(*a, **kw): + return results + + async def fake_ddg(*a, **kw): + return [] + + async def fake_rerank(q, c): + return [(u, t, 0.8) for u, t in c] + + monkeypatch.setattr("axon.routers.search.search_brave", fake_brave) + monkeypatch.setattr("axon.routers.search.search_ddg", fake_ddg) + monkeypatch.setattr("axon.routers.search.rerank_results", fake_rerank) + + from axon.routers.search import router + app = FastAPI() + app.include_router(router) + c = TestClient(app) + resp = c.post("/search", json={"query": "test", "domains_whitelist": ["allowed.com"]}) + assert resp.status_code == 200 + urls = [r["url"] for r in resp.json()["results"]] + assert all("allowed.com" in u for u in urls) + assert len(urls) == 1