gauravmeena0708/epfo-circulars

1import unittest2 3try:4    import pymupdf as fitz5except ImportError:6    import fitz7 8from document_assistant import (9    DocumentExtractionError,10    extract_pdf_text,11    format_conversation_history,12    select_document_context,13    uploaded_file_signature,14)15 16 17def make_pdf(text=""):18    document = fitz.open()19    page = document.new_page()20    if text:21        page.insert_textbox(fitz.Rect(36, 36, 560, 780), text, fontsize=10)22    data = document.tobytes()23    document.close()24    return data25 26 27class FakeOcrReader:28    def readtext(self, image, detail=0, paragraph=True):29        return ["Text recognized from a scanned noting-sheet page."]30 31 32class DocumentAssistantTests(unittest.TestCase):33    def test_upload_signature_changes_for_same_filename_with_new_content(self):34        first = uploaded_file_signature("note.pdf", b"first")35        second = uploaded_file_signature("note.pdf", b"second")36        self.assertNotEqual(first, second)37 38    def test_native_text_pdf_does_not_initialize_ocr(self):39        native_text = " ".join(f"word{number}" for number in range(40))40 41        def unexpected_ocr_load():42            raise AssertionError("OCR should not be loaded for a text page")43 44        result = extract_pdf_text(45            make_pdf(native_text),46            ocr_reader_factory=unexpected_ocr_load,47            native_text_min_words=25,48        )49 50        self.assertEqual(result.page_count, 1)51        self.assertEqual(result.scanned_page_count, 0)52        self.assertIn("word20", result.text)53 54    def test_short_native_text_is_preserved_when_ocr_is_unavailable(self):55        def unavailable_ocr():56            raise RuntimeError("model unavailable")57 58        result = extract_pdf_text(59            make_pdf("Short but important approval note"),60            ocr_reader_factory=unavailable_ocr,61            native_text_min_words=25,62        )63 64        self.assertIn("Short but important approval note", result.text)65        self.assertTrue(any("OCR could not be initialized" in item for item in result.warnings))66 67    def test_scanned_page_uses_ocr_reader(self):68        result = extract_pdf_text(69            make_pdf(),70            ocr_reader_factory=FakeOcrReader,71            native_text_min_words=25,72        )73 74        self.assertIn("Text recognized from a scanned", result.text)75        self.assertIn("Scanned Page", result.text)76 77    def test_invalid_pdf_has_safe_error(self):78        with self.assertRaisesRegex(DocumentExtractionError, "not a readable PDF"):79            extract_pdf_text(b"not a pdf")80 81    def test_long_document_context_keeps_relevant_and_boundary_pages(self):82        blocks = []83        for page_number in range(1, 13):84            body = f"ordinary material page {page_number} " * 2085            if page_number == 7:86                body += " unique_financial_sanction approval " * 1087            blocks.append(f"--- [Page {page_number} / Note Page] ---\n{body}")88        context, limited = select_document_context(89            "\n\n".join(blocks),90            "Explain the unique_financial_sanction",91            2_000,92        )93 94        self.assertTrue(limited)95        self.assertIn("Page 1 /", context)96        self.assertIn("Page 7 /", context)97        self.assertIn("Page 12 /", context)98 99    def test_conversation_history_uses_recent_turns(self):100        history = [101            {"role": "user", "content": "old question"},102            {"role": "assistant", "content": "old answer"},103            {"role": "user", "content": "recent follow-up"},104        ]105        formatted = format_conversation_history(history, max_chars=2_000, max_messages=2)106 107        self.assertNotIn("old question", formatted)108        self.assertIn("Assistant: old answer", formatted)109        self.assertIn("User: recent follow-up", formatted)110 111 112if __name__ == "__main__":113    unittest.main()114