backend/tests/test_ask.py

"""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