Spaces:
Running
Running
| """Regression tests for the issues found during the project review.""" | |
| import json | |
| import time | |
| from types import SimpleNamespace | |
| import fitz | |
| import pytest | |
| from fastapi.testclient import TestClient | |
| import main | |
| import backend.api.routes as routes | |
| from backend.core import settings | |
| from backend.core.config import Settings | |
| from backend.services import pdf_processor | |
| from backend.services.pdf_processor import PDFProcessor | |
| from backend.services.rag_engine import RAGEngine, select_cited_chunks | |
| from backend.services.rag_graph import DocuMindRAGGraph, format_history | |
| client = TestClient(main.app) | |
| # ── helpers ──────────────────────────────────────────────────────────────────── | |
| def make_pdf(path, texts): | |
| pdf = fitz.open() | |
| for text in texts: | |
| page = pdf.new_page() | |
| if text: | |
| page.insert_text((72, 72), text) | |
| pdf.save(path) | |
| pdf.close() | |
| return path.read_bytes() | |
| class FakeStore: | |
| def __init__(self): | |
| self.docs = {} | |
| self.queries = [] | |
| self.results = [] | |
| def list_documents(self, owner_id=None): | |
| return list(self.docs.values()) | |
| def document_exists(self, doc_id, owner_id=None): | |
| return doc_id in self.docs | |
| def index_document(self, processed, owner_id=""): | |
| self.docs[processed.doc_id] = { | |
| "doc_id": processed.doc_id, "filename": processed.filename, | |
| "total_pages": processed.total_pages, "total_chunks": processed.total_chunks, | |
| "file_size_kb": processed.file_size_kb, | |
| } | |
| return True | |
| def search(self, query, doc_ids=None, top_k=None, owner_id=None): | |
| self.queries.append(query) | |
| return self.results | |
| def fake_store(tmp_path, monkeypatch): | |
| store = FakeStore() | |
| monkeypatch.setattr(routes, "get_vector_store", lambda: store) | |
| monkeypatch.setattr(settings, "upload_folder", str(tmp_path / "uploads")) | |
| main._hits.clear() | |
| yield store | |
| main._hits.clear() | |
| # ── config ───────────────────────────────────────────────────────────────────── | |
| def test_placeholder_api_keys_are_treated_as_unset(): | |
| assert Settings(GROQ_API_KEY="replace_with_your_key").GROQ_API_KEY == "" | |
| assert Settings(api_key="replace_with_a_secret").api_key == "" | |
| assert Settings(GROQ_API_KEY=" gsk_real ").GROQ_API_KEY == "gsk_real" | |
| def test_default_chunk_size_fits_embedding_model_limit(): | |
| # all-MiniLM-L6-v2 truncates input beyond 256 tokens | |
| assert Settings().chunk_size <= 256 | |
| # ── PDF processing ───────────────────────────────────────────────────────────── | |
| def test_total_pages_counts_blank_trailing_pages(tmp_path): | |
| path = tmp_path / "two.pdf" | |
| make_pdf(path, ["This first page has plenty of meaningful text for a chunk.", None]) | |
| result = PDFProcessor().process(str(path)) | |
| assert result.total_pages == 2 | |
| def test_ocr_failure_gives_value_error_not_crash(tmp_path, monkeypatch): | |
| path = tmp_path / "scan.pdf" | |
| make_pdf(path, [None]) | |
| def broken_ocr(self, page): | |
| raise OSError("tesseract is not installed") # TesseractNotFoundError is an OSError | |
| monkeypatch.setattr(PDFProcessor, "_ocr_page", broken_ocr) | |
| with pytest.raises(ValueError, match="No extractable text"): | |
| PDFProcessor().process(str(path)) | |
| def test_chunking_works_when_tiktoken_is_unavailable(tmp_path, monkeypatch): | |
| def offline(_name): | |
| raise RuntimeError("no network") | |
| monkeypatch.setattr(pdf_processor.tiktoken, "get_encoding", offline) | |
| path = tmp_path / "doc.pdf" | |
| make_pdf(path, ["Offline chunking must still work and keep the page number intact."]) | |
| result = PDFProcessor().process(str(path)) | |
| assert result.chunks and result.chunks[0].page_number == 1 | |
| # ── grounding / citations ────────────────────────────────────────────────────── | |
| def _graph(): | |
| graph = object.__new__(DocuMindRAGGraph) | |
| graph.llm = None | |
| return graph | |
| def test_refusal_only_checked_at_start_of_answer(): | |
| chunks = [{"text": "Menu navigation settings profile options accounts dashboard"}] | |
| good = _graph().check_evidence({ | |
| "answer": "Navigation settings profile options accounts appear in the dashboard menu. " | |
| "Users cannot find advanced billing there.", | |
| "chunks": chunks}) | |
| refusal = _graph().check_evidence({ | |
| "answer": "I cannot find that information in the document.", "chunks": chunks}) | |
| assert good["is_grounded"] is True | |
| assert refusal["is_grounded"] is False | |
| def test_select_cited_chunks_uses_cited_sources(): | |
| chunks = [{"id": 1}, {"id": 2}, {"id": 3}] | |
| assert select_cited_chunks("Yes [Source 2].", chunks) == [{"id": 2}] | |
| def test_select_cited_chunks_falls_back_when_citation_is_invalid_or_missing(): | |
| chunks = [{"id": 1}, {"id": 2}] | |
| assert select_cited_chunks("Yes [Source 9].", chunks) == chunks | |
| assert select_cited_chunks("No citation at all.", chunks) == chunks | |
| def test_confidence_uses_all_retrieved_chunks_but_sources_use_cited(): | |
| fake_graph = SimpleNamespace(invoke=lambda **k: { | |
| "answer": "Weak answer [Source 2].", "is_grounded": True, | |
| "chunks": [{"doc_name": "a.pdf", "page_number": 1, "text": "x" * 10, "similarity": 0.95}, | |
| {"doc_name": "a.pdf", "page_number": 2, "text": "y" * 10, "similarity": 0.80}]}) | |
| engine = object.__new__(RAGEngine) | |
| engine.graph = fake_graph | |
| result = engine.query("question here") | |
| assert [s.page_number for s in result.sources] == [2] | |
| assert result.confidence == "HIGH" # from top similarity 0.95, not just the cited chunk | |
| def test_format_history_accepts_dicts_and_truncates(): | |
| text = format_history([{"role": "user", "content": "a" * 5000}, {"role": "bogus", "content": "b"}]) | |
| assert text.splitlines()[0].startswith("user: ") and len(text.splitlines()[0]) <= 2006 | |
| assert text.splitlines()[1] == "user: b" # unknown roles are neutralised | |
| # ── request validation ───────────────────────────────────────────────────────── | |
| def test_invalid_history_is_rejected(history, fake_store): | |
| response = client.post("/api/query", json={"question": "valid question", "history": history}) | |
| assert response.status_code == 422 | |
| # ── upload behaviour ─────────────────────────────────────────────────────────── | |
| def test_duplicate_upload_reports_already_indexed(tmp_path, fake_store): | |
| content = make_pdf(tmp_path / "x.pdf", ["Duplicate detection needs enough readable text here."]) | |
| files = {"file": ("x.pdf", content, "application/pdf")} | |
| first = client.post("/api/upload", files=files) | |
| second = client.post("/api/upload", files=files) | |
| assert first.status_code == second.status_code == 200 | |
| assert "uploaded successfully" in first.json()["message"] | |
| assert second.json()["message"] == "Document already indexed." | |
| assert first.json()["doc_id"] == second.json()["doc_id"] | |
| assert len(fake_store.docs) == 1 | |
| def test_oversize_upload_is_rejected(monkeypatch, fake_store): | |
| monkeypatch.setattr(routes, "MAX_UPLOAD", 1024) | |
| response = client.post("/api/upload", files={"file": ("big.pdf", b"%PDF-" + b"0" * 4096, "application/pdf")}) | |
| assert response.status_code == 400 | |
| assert "too large" in response.json()["detail"].lower() | |
| def test_large_upload_runs_as_background_job(tmp_path, monkeypatch, fake_store): | |
| monkeypatch.setattr(routes, "BACKGROUND_THRESHOLD", 1) | |
| content = make_pdf(tmp_path / "bg.pdf", ["Background ingestion needs some readable text as well."]) | |
| response = client.post("/api/upload", files={"file": ("bg.pdf", content, "application/pdf")}) | |
| body = response.json() | |
| assert response.status_code == 200 and body["job_id"] and body["status"] == "queued" | |
| status = client.get(f"/api/jobs/{body['job_id']}").json() | |
| assert status["status"] == "done" and status["doc_id"] | |
| assert "owner_id" not in status and "created" not in status | |
| assert client.get("/api/jobs/does-not-exist").status_code == 404 | |
| def test_finished_jobs_are_pruned_after_ttl(): | |
| routes.jobs["old"] = {"status": "done", "created": time.time() - routes.JOB_TTL_SECONDS - 5} | |
| routes.jobs["fresh"] = {"status": "done", "created": time.time()} | |
| routes.jobs["running"] = {"status": "processing", "created": 0} | |
| routes._prune_jobs() | |
| assert "old" not in routes.jobs | |
| assert "fresh" in routes.jobs and "running" in routes.jobs | |
| routes.jobs.pop("fresh", None); routes.jobs.pop("running", None) | |
| # ── middleware: rate limit + request id + safe errors ────────────────────────── | |
| def test_query_rate_limit_returns_429(fake_store): | |
| codes = [client.post("/api/query", json={"question": "a valid question"}).status_code | |
| for _ in range(31)] | |
| assert 429 not in codes[:30] and codes[30] == 429 | |
| def test_request_id_is_sanitised(fake_store): | |
| good = client.get("/api/health", headers={"X-Request-ID": "abc-123_ok"}) | |
| bad = client.get("/api/health", headers={"X-Request-ID": "bad id\twith <script>"}) | |
| assert good.headers["X-Request-ID"] == "abc-123_ok" | |
| assert "<" not in bad.headers["X-Request-ID"] and " " not in bad.headers["X-Request-ID"] | |
| def test_storage_failure_is_a_generic_500(monkeypatch): | |
| class Broken: | |
| def list_documents(self, owner_id=None): | |
| raise RuntimeError("secret path /var/db exploded") | |
| monkeypatch.setattr(routes, "get_vector_store", lambda: Broken()) | |
| safe_client = TestClient(main.app, raise_server_exceptions=False) | |
| response = safe_client.get("/api/documents") | |
| assert response.status_code == 500 | |
| assert "secret" not in response.text and "exploded" not in response.text | |
| # ── streaming endpoint ───────────────────────────────────────────────────────── | |
| def _events(response): | |
| events, current = [], {} | |
| for line in response.text.splitlines(): | |
| if line.startswith("event: "): | |
| current = {"event": line[7:]} | |
| elif line.startswith("data: "): | |
| current["data"] = json.loads(line[6:]) | |
| events.append(current) | |
| return events | |
| def _fake_engine(graph): | |
| return SimpleNamespace(graph=graph, _assess_confidence=RAGEngine._assess_confidence) | |
| def test_stream_no_evidence_path(fake_store, monkeypatch): | |
| fake_store.docs["d1"] = {"doc_id": "d1", "filename": "a.pdf", "total_pages": 1, | |
| "total_chunks": 1, "file_size_kb": 1.0} | |
| fake_store.results = [] | |
| monkeypatch.setattr(routes, "get_rag_engine", lambda: _fake_engine(SimpleNamespace(llm=None))) | |
| response = client.post("/api/query/stream", json={"question": "anything here?"}) | |
| final = _events(response)[-1] | |
| assert final["event"] == "final" | |
| assert final["data"]["confidence"] == "LOW" and final["data"]["sources"] == [] | |
| def test_stream_uses_condensed_question_for_retrieval_and_filters_sources(fake_store, monkeypatch): | |
| fake_store.docs["d1"] = {"doc_id": "d1", "filename": "a.pdf", "total_pages": 2, | |
| "total_chunks": 2, "file_size_kb": 1.0} | |
| fake_store.results = [ | |
| {"doc_name": "a.pdf", "page_number": 1, "text": "Revenue grew twelve percent.", "similarity": 0.9}, | |
| {"doc_name": "a.pdf", "page_number": 2, "text": "Costs fell slightly.", "similarity": 0.7}, | |
| ] | |
| llm = SimpleNamespace(stream=lambda messages: [ | |
| SimpleNamespace(content="Revenue grew "), SimpleNamespace(content="twelve percent [Source 1].")]) | |
| graph = SimpleNamespace( | |
| llm=llm, | |
| condense_question=lambda state: {"question": "What happened to revenue in 2024?"}, | |
| check_evidence=lambda state: {"is_grounded": True}) | |
| monkeypatch.setattr(routes, "get_rag_engine", lambda: _fake_engine(graph)) | |
| response = client.post("/api/query/stream", json={ | |
| "question": "and what about it?", | |
| "history": [{"role": "user", "content": "Tell me about revenue"}]}) | |
| events = _events(response) | |
| assert fake_store.queries == ["What happened to revenue in 2024?"] # condensed BEFORE search | |
| assert [e["event"] for e in events] == ["token", "token", "final"] | |
| final = events[-1]["data"] | |
| assert final["answer"] == "Revenue grew twelve percent [Source 1]." | |
| assert [s["page_number"] for s in final["sources"]] == [1] | |
| assert final["confidence"] == "HIGH" and final["is_grounded"] is True | |