gauravmeena0708/epfo-circulars

1import io2import unittest3from PIL import Image4import fitz5 6import pdf_utils7 8 9def _create_sample_pdf(page_count: int = 3, text_prefix: str = "Page") -> bytes:10    """Helper to create an in-memory PDF with specified number of pages."""11    doc = fitz.open()12    for i in range(page_count):13        page = doc.new_page(width=595, height=842)14        page.insert_text((50, 72), f"{text_prefix} {i + 1} content", fontsize=14)15    pdf_bytes = doc.tobytes()16    doc.close()17    return pdf_bytes18 19 20def _create_sample_image(width: int = 100, height: int = 100, color: str = "red") -> bytes:21    """Helper to create an in-memory PNG image."""22    img = Image.new("RGB", (width, height), color=color)23    buf = io.BytesIO()24    img.save(buf, format="PNG")25    return buf.getvalue()26 27 28class TestPDFUtils(unittest.TestCase):29 30    def test_merge_pdfs(self):31        pdf1 = _create_sample_pdf(2, "Doc1 Page")32        pdf2 = _create_sample_pdf(3, "Doc2 Page")33        pdf3 = _create_sample_pdf(1, "Doc3 Page")34 35        merged = pdf_utils.merge_pdfs([pdf1, pdf2, pdf3])36        doc = fitz.open(stream=merged, filetype="pdf")37        self.assertEqual(len(doc), 6)38        doc.close()39 40    def test_merge_pdfs_empty_raises(self):41        with self.assertRaises(ValueError):42            pdf_utils.merge_pdfs([])43 44    def test_merge_pdfs_corrupted_raises(self):45        with self.assertRaises(ValueError):46            pdf_utils.merge_pdfs([b"not a valid pdf", _create_sample_pdf(1)])47 48    def test_parse_page_ranges(self):49        self.assertEqual(pdf_utils.parse_page_ranges("1", 5), [0])50        self.assertEqual(pdf_utils.parse_page_ranges("1-3", 5), [0, 1, 2])51        self.assertEqual(pdf_utils.parse_page_ranges("1, 3, 5", 5), [0, 2, 4])52        self.assertEqual(pdf_utils.parse_page_ranges("2-4, 1", 5), [1, 2, 3, 0])53        self.assertEqual(pdf_utils.parse_page_ranges("all", 3), [0, 1, 2])54 55    def test_parse_page_ranges_invalid(self):56        with self.assertRaises(ValueError):57            pdf_utils.parse_page_ranges("0", 5)  # 1-indexed58        with self.assertRaises(ValueError):59            pdf_utils.parse_page_ranges("6", 5)  # Out of bounds60        with self.assertRaises(ValueError):61            pdf_utils.parse_page_ranges("abc", 5)62        with self.assertRaises(ValueError):63            pdf_utils.parse_page_ranges("5-2", 5)  # Reversed range64 65    def test_split_pdf(self):66        pdf = _create_sample_pdf(5, "Page")67        split_bytes = pdf_utils.split_pdf(pdf, "2-4")68        doc = fitz.open(stream=split_bytes, filetype="pdf")69        self.assertEqual(len(doc), 3)70        doc.close()71 72    def test_split_pdf_to_individual_pages(self):73        pdf = _create_sample_pdf(3, "Page")74        pages = pdf_utils.split_pdf_to_individual_pages(pdf)75        self.assertEqual(len(pages), 3)76        self.assertEqual(pages[0][0], "page_1.pdf")77        doc = fitz.open(stream=pages[0][1], filetype="pdf")78        self.assertEqual(len(doc), 1)79        doc.close()80 81    def test_rotate_pdf_pages(self):82        pdf = _create_sample_pdf(3, "Page")83        rotated = pdf_utils.rotate_pdf_pages(pdf, 90)84        doc = fitz.open(stream=rotated, filetype="pdf")85        for page in doc:86            self.assertEqual(page.rotation, 90)87        doc.close()88 89    def test_rotate_specific_pages(self):90        pdf = _create_sample_pdf(3, "Page")91        rotated = pdf_utils.rotate_pdf_pages(pdf, 180, page_range_str="2")92        doc = fitz.open(stream=rotated, filetype="pdf")93        self.assertEqual(doc[0].rotation, 0)94        self.assertEqual(doc[1].rotation, 180)95        self.assertEqual(doc[2].rotation, 0)96        doc.close()97 98    def test_rotate_invalid_degree(self):99        pdf = _create_sample_pdf(1)100        with self.assertRaises(ValueError):101            pdf_utils.rotate_pdf_pages(pdf, 45)102 103    def test_compress_pdf(self):104        pdf = _create_sample_pdf(3, "Page")105        compressed_bytes, orig_size, new_size = pdf_utils.compress_pdf(pdf)106        self.assertGreater(len(compressed_bytes), 0)107        self.assertEqual(orig_size, len(pdf))108        self.assertEqual(new_size, len(compressed_bytes))109        doc = fitz.open(stream=compressed_bytes, filetype="pdf")110        self.assertEqual(len(doc), 3)111        doc.close()112 113    def test_images_to_pdf(self):114        img1 = _create_sample_image(100, 100, "red")115        img2 = _create_sample_image(120, 80, "blue")116        pdf_bytes = pdf_utils.images_to_pdf([("doc1.png", img1), ("doc2.png", img2)])117        doc = fitz.open(stream=pdf_bytes, filetype="pdf")118        self.assertEqual(len(doc), 2)119        doc.close()120 121    def test_pdf_to_images(self):122        pdf = _create_sample_pdf(2, "Page")123        images = pdf_utils.pdf_to_images(pdf, dpi=100, image_format="PNG")124        self.assertEqual(len(images), 2)125        self.assertEqual(images[0][0], "page_1.png")126        self.assertEqual(images[1][0], "page_2.png")127        # Verify it can be loaded by PIL128        img = Image.open(io.BytesIO(images[0][1]))129        self.assertGreater(img.width, 0)130 131 132if __name__ == "__main__":133    unittest.main()134