arsam
6e4ffcca7a
Import the existing local project snapshot. This commit organizes existing files; it does not represent newly developed functionality.
70 lines
2.2 KiB
Python
70 lines
2.2 KiB
Python
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
|
|
def test_ingest_and_search_private_index(tmp_path, monkeypatch):
|
|
monkeypatch.setattr("axon.services.index.settings.index_db_path", str(tmp_path / "index.sqlite3"))
|
|
|
|
from axon.routers.index import router
|
|
|
|
app = FastAPI()
|
|
app.include_router(router)
|
|
client = TestClient(app)
|
|
|
|
ingest = client.post(
|
|
"/ingest",
|
|
json={
|
|
"documents": [
|
|
{
|
|
"id": "doc1",
|
|
"title": "Private Axon Notes",
|
|
"content": "Axon supports private retrieval infrastructure.",
|
|
"url": "private://doc1",
|
|
}
|
|
]
|
|
},
|
|
)
|
|
assert ingest.status_code == 200
|
|
assert ingest.json()["indexed"] == 1
|
|
|
|
search = client.post("/index/search", json={"query": "private retrieval"})
|
|
assert search.status_code == 200
|
|
data = search.json()
|
|
assert data["results"][0]["id"] == "doc1"
|
|
|
|
|
|
def test_search_can_include_private_index(monkeypatch):
|
|
from axon.models import SearchResult
|
|
from axon.services.cache import InMemoryCache
|
|
|
|
cache = InMemoryCache()
|
|
|
|
async def async_cache():
|
|
return cache
|
|
|
|
async def empty(*args, **kwargs):
|
|
return []
|
|
|
|
async def private(*args, **kwargs):
|
|
return [SearchResult(title="Private", url="private://1", content="secret axon data", score=0.9)]
|
|
|
|
async def rerank(query, contents):
|
|
return [(url, text, 0.9) for url, text in contents]
|
|
|
|
monkeypatch.setattr("axon.routers.search.get_cache", async_cache)
|
|
monkeypatch.setattr("axon.routers.search.search_brave", empty)
|
|
monkeypatch.setattr("axon.routers.search.search_ddg", empty)
|
|
monkeypatch.setattr("axon.routers.search.search_serper", empty)
|
|
monkeypatch.setattr("axon.routers.search.search_index_as_results", private)
|
|
monkeypatch.setattr("axon.routers.search.rerank_results", rerank)
|
|
|
|
from axon.routers.search import router
|
|
|
|
app = FastAPI()
|
|
app.include_router(router)
|
|
client = TestClient(app)
|
|
|
|
resp = client.post("/search", json={"query": "secret"})
|
|
assert resp.status_code == 200
|
|
assert resp.json()["results"][0]["url"] == "private://1"
|