arsam
c71246da26
Import the existing local project snapshot. This commit organizes existing files; it does not represent newly developed functionality.
93 lines
3.0 KiB
Python
93 lines
3.0 KiB
Python
import pytest
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
from axon.models import SearchResponse, SearchResult
|
|
|
|
|
|
def _result(url: str, title: str = "Title") -> SearchResult:
|
|
return SearchResult(title=title, url=url, content="content", score=0.9)
|
|
|
|
|
|
@pytest.fixture
|
|
def client(monkeypatch):
|
|
async def fake_generate(query, count=3):
|
|
return [f"{query} angle {i}" for i in range(count)]
|
|
|
|
async def fake_search(request):
|
|
return SearchResponse(
|
|
query=request.query,
|
|
results=[
|
|
_result("https://example.com/shared", "Shared"),
|
|
_result(f"https://example.com/{request.query.replace(' ', '-')}"),
|
|
],
|
|
response_time=0.1,
|
|
)
|
|
|
|
async def fake_report(query, results, output_format="markdown"):
|
|
return f"report:{query}:{len(results)}:{output_format}"
|
|
|
|
monkeypatch.setattr("axon.routers.research.generate_sub_queries", fake_generate)
|
|
monkeypatch.setattr("axon.routers.research.search", fake_search)
|
|
monkeypatch.setattr("axon.routers.research.synthesize_report", fake_report)
|
|
|
|
from axon.routers.research import router
|
|
|
|
app = FastAPI()
|
|
app.include_router(router)
|
|
return TestClient(app)
|
|
|
|
|
|
def test_research_basic_runs_three_sub_queries(client):
|
|
resp = client.post("/research", json={"query": "axon"})
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert len(data["sub_queries"]) == 3
|
|
assert data["report"] == "report:axon:4:markdown"
|
|
assert len(data["sources"]) == 4
|
|
|
|
|
|
def test_research_deep_runs_five_sub_queries(client):
|
|
resp = client.post("/research", json={"query": "axon", "depth": "deep"})
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert len(data["sub_queries"]) == 5
|
|
assert len(data["sources"]) == 6
|
|
|
|
|
|
def test_research_passes_json_output_format(client):
|
|
resp = client.post("/research", json={"query": "axon", "output_format": "json"})
|
|
assert resp.status_code == 200
|
|
assert resp.json()["report"].endswith(":json")
|
|
|
|
|
|
def test_research_skips_failed_sub_search(monkeypatch):
|
|
async def fake_generate(query, count=3):
|
|
return ["one", "two", "three"]
|
|
|
|
async def fake_search(request):
|
|
if request.query == "two":
|
|
raise RuntimeError("source failed")
|
|
return SearchResponse(
|
|
query=request.query,
|
|
results=[_result(f"https://example.com/{request.query}")],
|
|
response_time=0.1,
|
|
)
|
|
|
|
async def fake_report(query, results, output_format="markdown"):
|
|
return f"count:{len(results)}"
|
|
|
|
monkeypatch.setattr("axon.routers.research.generate_sub_queries", fake_generate)
|
|
monkeypatch.setattr("axon.routers.research.search", fake_search)
|
|
monkeypatch.setattr("axon.routers.research.synthesize_report", fake_report)
|
|
|
|
from axon.routers.research import router
|
|
|
|
app = FastAPI()
|
|
app.include_router(router)
|
|
client = TestClient(app)
|
|
|
|
resp = client.post("/research", json={"query": "axon"})
|
|
assert resp.status_code == 200
|
|
assert resp.json()["report"] == "count:2"
|