arsam
240aacc030
Import the existing local project snapshot. This commit organizes existing files; it does not represent newly developed functionality.
126 lines
3.6 KiB
Python
126 lines
3.6 KiB
Python
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
|