File size: 1,748 Bytes
d9f07b6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | """Tests for the retrieval + chat loop (Session 3)."""
from fastapi.testclient import TestClient
from app.main import app
from app.retrieval import DocumentIndex, best_sentence
client = TestClient(app)
FACTS = [
"The mitochondria is the powerhouse of the cell.",
"Python is a popular programming language for data science.",
"The Eiffel Tower is located in Paris, France.",
]
def _ingest(text: str) -> str:
response = client.post("/documents", data={"text": text})
assert response.status_code == 201
return response.json()["document_id"]
def test_index_ranks_relevant_chunk_first():
index = DocumentIndex(FACTS)
hits = index.search("Where is the Eiffel Tower?", k=1)
assert FACTS[hits[0][0]] == FACTS[2]
def test_best_sentence_picks_overlapping_sentence():
passage = " ".join(FACTS)
assert "Eiffel" in best_sentence(passage, "Where is the Eiffel Tower?")
def test_ask_returns_relevant_passage_and_answer():
doc_id = _ingest(" ".join(FACTS))
response = client.post(
"/ask",
json={"document_id": doc_id, "question": "Where is the Eiffel Tower?"},
)
assert response.status_code == 200
body = response.json()
assert "Paris" in body["source_passage"]
assert "Eiffel" in body["answer"] or "Paris" in body["answer"]
assert body["chunk_index"] >= 0
def test_ask_empty_question_is_rejected():
doc_id = _ingest(" ".join(FACTS))
response = client.post("/ask", json={"document_id": doc_id, "question": " "})
assert response.status_code == 400
def test_ask_unknown_document_returns_404():
response = client.post(
"/ask", json={"document_id": "does-not-exist", "question": "Hello?"}
)
assert response.status_code == 404
|