arsam
41033ad099
Import the existing local project snapshot. This commit organizes existing files; it does not represent newly developed functionality.
85 lines
2.3 KiB
Python
85 lines
2.3 KiB
Python
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
|
|
def test_citation_validation_route():
|
|
from axon.routers.validate import router
|
|
|
|
app = FastAPI()
|
|
app.include_router(router)
|
|
client = TestClient(app)
|
|
|
|
resp = client.post(
|
|
"/validate/citations",
|
|
json={
|
|
"report": "Claim [1] and missing [3]",
|
|
"sources": [
|
|
{"title": "A", "url": "https://a.com", "content": "a"},
|
|
{"title": "B", "url": "https://b.com", "content": "b"},
|
|
],
|
|
},
|
|
)
|
|
assert resp.status_code == 200
|
|
assert resp.json()["missing_citations"] == [3]
|
|
assert resp.json()["valid"] is False
|
|
|
|
|
|
def test_tavily_search_adapter(monkeypatch):
|
|
from axon.models import SearchResponse
|
|
|
|
async def fake_search(request):
|
|
assert request.domains_whitelist == ["example.com"]
|
|
return SearchResponse(query=request.query, results=[], response_time=0.1)
|
|
|
|
monkeypatch.setattr("axon.routers.tavily.search", fake_search)
|
|
|
|
from axon.routers.tavily import router
|
|
|
|
app = FastAPI()
|
|
app.include_router(router)
|
|
client = TestClient(app)
|
|
|
|
resp = client.post(
|
|
"/tavily/search",
|
|
json={"query": "axon", "include_domains": ["example.com"]},
|
|
)
|
|
assert resp.status_code == 200
|
|
assert resp.json()["query"] == "axon"
|
|
|
|
|
|
def test_job_lifecycle(monkeypatch):
|
|
from axon.models import ResearchResponse
|
|
|
|
async def fake_research(request):
|
|
return ResearchResponse(
|
|
query=request.query,
|
|
report="done",
|
|
sources=[],
|
|
sub_queries=[],
|
|
response_time=0.1,
|
|
)
|
|
|
|
monkeypatch.setattr("axon.routers.jobs.research", fake_research)
|
|
|
|
from axon.routers.jobs import router
|
|
|
|
app = FastAPI()
|
|
app.include_router(router)
|
|
client = TestClient(app)
|
|
|
|
create = client.post("/research/jobs", json={"query": "axon"})
|
|
assert create.status_code == 200
|
|
job_id = create.json()["job_id"]
|
|
|
|
import time
|
|
|
|
for _ in range(20):
|
|
read = client.get(f"/jobs/{job_id}")
|
|
if read.json()["status"] == "succeeded":
|
|
break
|
|
time.sleep(0.01)
|
|
|
|
assert read.status_code == 200
|
|
assert read.json()["status"] == "succeeded"
|
|
assert read.json()["result"]["report"] == "done"
|