import sys from unittest.mock import MagicMock import multiprocessing # Prevent downloading/loading Hugging Face embeddings during testing by mocking the class import langchain_community.embeddings langchain_community.embeddings.HuggingFaceEmbeddings = MagicMock() import pytest from fastapi.testclient import TestClient from main import ( app, detect_question_intent, sanitize_upload_filename, concise_excerpt, split_sentences, clean_sentence, query_keywords, tokenize_text, build_answer_from_documents, INSUFFICIENT_CONTEXT_MESSAGE, passes_evidence_gate, document_dedupe_key, citation_source_for_document, internal_token_valid, normalize_session_id, get_session_dir, _extract_pdf_text_worker, ) import secrets as _secrets def is_authorized_session_update(session: dict, provided_secret) -> bool: """Replicate the session-secret check from the endpoint (moved inline upstream).""" expected = (session.get("session_secret") or "").strip() candidate = (provided_secret or "").strip() if not expected or not candidate: return False return _secrets.compare_digest(candidate, expected) def test_session_secret_authorizes_only_matching_secret(): session = {"session_secret": "expected-secret"} assert is_authorized_session_update(session, "expected-secret") is True assert is_authorized_session_update(session, "wrong-secret") is False assert is_authorized_session_update(session, None) is False assert is_authorized_session_update({}, "expected-secret") is False def test_detect_question_intent(): assert detect_question_intent("What is this document about?") == "overview" assert detect_question_intent("What are these documents about?") == "overview" assert detect_question_intent("Explain the connection between X and Y") == "relationship" assert detect_question_intent("How does X compare to Y?") == "comparison" assert detect_question_intent("Compare the performance of model A and B") == "comparison" assert detect_question_intent("What is the revenue in 2023?") == "factual" assert detect_question_intent("Who is the CEO of the company?") == "factual" def test_sanitize_upload_filename_valid(): assert sanitize_upload_filename("test.pdf") == "test.pdf" assert sanitize_upload_filename("path/to/my_document.PDF") == "my_document.PDF" assert sanitize_upload_filename("C:\\Users\\file-name_123.pdf") == "file-name_123.pdf" def test_sanitize_upload_filename_invalid(): with pytest.raises(ValueError, match="Missing PDF file path"): sanitize_upload_filename("") with pytest.raises(ValueError, match="Missing PDF file path"): sanitize_upload_filename(" ") with pytest.raises(ValueError, match="Only PDF files are allowed"): sanitize_upload_filename("test.txt") with pytest.raises(ValueError, match="Uploaded filename contains unsupported characters"): sanitize_upload_filename("test$file.pdf") def test_internal_token_valid_allows_when_unset(): assert internal_token_valid(None, "") is True assert internal_token_valid("", "") is True def test_internal_token_valid_rejects_missing_when_set(): assert internal_token_valid(None, "secret") is False assert internal_token_valid("", "secret") is False assert internal_token_valid(" ", "secret") is False def test_internal_token_valid_accepts_exact_match(): assert internal_token_valid("secret", "secret") is True def test_internal_auth_middleware_protects_validate_session_write(): import main as main_module original_token = main_module.INTERNAL_RAG_TOKEN main_module.INTERNAL_RAG_TOKEN = "test-secret" try: client = TestClient(app) response = client.post("/validate-session-write") assert response.status_code == 403 assert response.json()["error"] == "Forbidden" finally: main_module.INTERNAL_RAG_TOKEN = original_token def test_normalize_session_id_rejects_invalid_values(): with pytest.raises(ValueError, match="Missing session id"): normalize_session_id("") with pytest.raises(ValueError): normalize_session_id("not-a-uuid") def test_get_session_dir_requires_uuid_session_id(): with pytest.raises(ValueError): get_session_dir("../escape") def test_normalize_session_id_returns_canonical_uuid(): normalized = normalize_session_id("550E8400-E29B-41D4-A716-446655440000") assert normalized == "550e8400-e29b-41d4-a716-446655440000" def test_extract_pdf_text_worker_enforces_page_limit(tmp_path): from pypdf import PdfWriter pdf_path = tmp_path / "hello.pdf" writer = PdfWriter() writer.add_blank_page(width=300, height=144) with pdf_path.open("wb") as fp: writer.write(fp) # Use a local queue and call the worker directly (no subprocess) to validate limit logic. q = multiprocessing.Queue(maxsize=1) _extract_pdf_text_worker(str(pdf_path), max_pages=0, max_chars=1000, out_queue=q) result = q.get(timeout=2) assert result["ok"] is False assert "too many pages" in result["error"].lower() def test_concise_excerpt(): text = "This is a very long sentence that we want to abbreviate cleanly." assert concise_excerpt(text, max_chars=20) == "This is a very long..." assert concise_excerpt(text, max_chars=100) == text def test_split_sentences(): text = "First sentence! Second sentence. Third one?" sentences = split_sentences(text) assert len(sentences) == 3 assert sentences[0] == "First sentence!" assert sentences[1] == "Second sentence." assert sentences[2] == "Third one?" def test_clean_sentence(): assert clean_sentence(" - Clean this sentence ") == "Clean this sentence" assert clean_sentence("* clean me ") == "clean me" def test_query_keywords(): # Stopwords like "what", "is", "this", "about" are filtered out # Only tokens with length > 2 are kept assert query_keywords("What is this document about revenue?") == {"revenue"} assert query_keywords("accuracy of model") == {"model", "accuracy"} def test_empty_query_handling(): query = "" assert query.strip() == "" def test_invalid_query_type(): query = None assert query is None def test_context_document_presence(): docs = ["sample pdf content", "rag pipeline notes"] assert len(docs) > 0 def test_answer_response_structure(): response = { "answer": "Sample answer", "sources": ["doc1.pdf"] } assert "answer" in response assert isinstance(response["sources"], list) def test_db_connection_placeholder(): db_status = True assert db_status is True def test_query_routing_logic(): query = "Summarize this PDF" if "summarize" in query.lower(): route = "summarizer" else: route = "qa" assert route == "summarizer" def test_crawler_response_structure(): crawler_output = { "url": "https://example.com", "content": "Sample crawled content", "status": 200 } assert "url" in crawler_output assert "content" in crawler_output assert crawler_output["status"] == 200 def test_crawler_empty_content(): crawler_output = { "url": "https://example.com", "content": "", "status": 200 } assert crawler_output["content"] == "" def test_crawler_failed_status(): crawler_output = { "url": "https://example.com", "content": None, "status": 500 } assert crawler_output["status"] >= 400 def test_retry_logic_placeholder(): retries = 3 success = True for _ in range(retries): success = True assert success is True def test_document_extraction_consistency(): extracted_chunks = [ "chunk one", "chunk two", "chunk three" ] assert len(extracted_chunks) == 3 assert all(isinstance(chunk, str) for chunk in extracted_chunks) def test_crawler_metadata_preservation(): metadata = { "source": "sample.pdf", "page": 1 } assert metadata["source"] == "sample.pdf" assert metadata["page"] == 1 def test_empty_document_handling(): extracted_text = "" assert extracted_text == "" def test_unstructured_pdf_ingestion_mock(): mock_document = { "filename": "research.pdf", "content": "This is extracted PDF content" } assert "pdf" in mock_document["filename"] assert len(mock_document["content"]) > 0 class DummyDocument: def __init__(self, content, filename="doc.pdf", page=0): self.page_content = content self.metadata = { "filename": filename, "page": page, "document_id": filename, } def test_evidence_gate_refuses_when_overlap_missing(): docs = [DummyDocument("this is unrelated content", filename="a.pdf", page=0)] assert passes_evidence_gate("What is the revenue?", docs, best_score=0.1, intent="factual") is False def test_evidence_gate_allows_when_overlap_and_score_good(): docs = [DummyDocument("Revenue for 2023 was 10 million.", filename="a.pdf", page=0)] assert passes_evidence_gate("What is the revenue for 2023?", docs, best_score=0.2, intent="factual") is True def test_build_answer_includes_citations_for_grounded_answer(): doc = DummyDocument("Revenue for 2023 was 10 million.", filename="a.pdf", page=0) source_id_by_key = {document_dedupe_key(doc): 1} answer = build_answer_from_documents( "What is the revenue for 2023?", [doc], "factual", source_id_by_key=source_id_by_key, ) assert "Source 1" in answer or "Sources 1" in answer def test_build_answer_refuses_when_unanswerable(): doc = DummyDocument("This document is about hiring policies.", filename="a.pdf", page=0) source_id_by_key = {document_dedupe_key(doc): 1} answer = build_answer_from_documents( "What is the revenue for 2023?", [doc], "factual", source_id_by_key=source_id_by_key, ) assert answer == INSUFFICIENT_CONTEXT_MESSAGE def test_citation_source_for_document_preserves_jump_metadata(): doc = DummyDocument("Internship duration is 6 weeks. More details follow.", filename="policy.pdf", page=11) doc.metadata["chunk_index"] = 4 source = citation_source_for_document(doc, 0) assert source["document"] == "policy.pdf" assert source["page"] == 12 assert source["chunk_index"] == 4 assert source["text"].startswith("Internship duration") assert source["preview"].startswith("Internship duration") def test_citation_source_for_document_handles_missing_metadata(): doc = DummyDocument("Useful supporting text.", filename="", page=0) doc.metadata = {} source = citation_source_for_document(doc, 2) assert source["document"] == "uploaded document" assert source["page"] is None assert source["chunk_index"] == 2