sanatan_ai / tests /test_entity_search_api.py
vikramvasudevan's picture
Upload folder using huggingface_hub
595350e verified
Raw
History Blame Contribute Delete
2.75 kB
import json
from fastapi.testclient import TestClient
from fastapi import FastAPI
from server import router
import logging
# Setup test app
app = FastAPI()
app.include_router(router, prefix="/api")
client = TestClient(app)
def print_results(title, data):
print(f"\n--- {title} ---")
print(f"Total Matches: {data.get('total_matches')}")
print(f"Page: {data.get('page')}, Page Size: {data.get('page_size')}")
results = data.get("results", [])
print(f"Results Count: {len(results)}")
for r in results:
print(f" - [{r['scripture_name']}] {r['relative_path']} (Index: {r['_global_index']})")
def test_topic_bhakti_pagination():
print("\nTesting Topic: Bhakti (Pagination)")
# Page 1
resp1 = client.post("/api/search/entity", json={
"entity_value": "Bhakti",
"entity_type": "topic",
"page": 1,
"page_size": 3
})
assert resp1.status_code == 200
data1 = resp1.json()
print_results("Bhakti - Page 1", data1)
if data1["total_matches"] > 3:
# Page 2
resp2 = client.post("/api/search/entity", json={
"entity_value": "Bhakti",
"entity_type": "topic",
"page": 2,
"page_size": 3
})
assert resp2.status_code == 200
data2 = resp2.json()
print_results("Bhakti - Page 2", data2)
# Verify pagination uniqueness
indices1 = [r["_global_index"] for r in data1["results"]]
indices2 = [r["_global_index"] for r in data2["results"]]
# Note: Uniqueness depends on stable ordering in DB, which we added (scripture_name, _global_index)
intersection = set(indices1).intersection(set(indices2))
# Simple check: if they are different pages, indices shouldn't overlap unless the data itself repeats across scriptures with same index
# But since we use stable ordering, they should be distinct segments.
print(f"Overlap check: {intersection}")
def test_character_prahlad():
print("\nTesting Character: Prahlad")
resp = client.post("/api/search/entity", json={
"entity_value": "Prahlad",
"entity_type": "character",
"page": 1,
"page_size": 5
})
assert resp.status_code == 200
data = resp.json()
print_results("Prahlad Results", data)
if __name__ == "__main__":
# Configure logging to be less verbose during tests
logging.getLogger("db").setLevel(logging.WARNING)
try:
test_topic_bhakti_pagination()
test_character_prahlad()
print("\nAll entity search tests completed successfully!")
except Exception as e:
print(f"\nTest failed: {e}")
import traceback
traceback.print_exc()