mirror of
https://github.com/denizsafak/abogen.git
synced 2026-07-18 05:40:26 +02:00
118 lines
4.0 KiB
Python
118 lines
4.0 KiB
Python
import unittest
|
|
import os
|
|
import shutil
|
|
import fitz # PyMuPDF
|
|
from abogen.book_parser import PdfParser
|
|
|
|
class TestPdfStructure(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
self.test_dir = "tests/test_data_pdf"
|
|
if os.path.exists(self.test_dir):
|
|
shutil.rmtree(self.test_dir)
|
|
os.makedirs(self.test_dir)
|
|
self.pdf_path = os.path.join(self.test_dir, "structure_test.pdf")
|
|
|
|
def tearDown(self):
|
|
if os.path.exists(self.test_dir):
|
|
shutil.rmtree(self.test_dir)
|
|
|
|
def test_pdf_structure_with_toc(self):
|
|
# Create PDF
|
|
doc = fitz.open()
|
|
p1 = doc.new_page()
|
|
p1.insert_text((50,50), "Page 1 Content")
|
|
p2 = doc.new_page()
|
|
p2.insert_text((50,50), "Page 2 Content")
|
|
p3 = doc.new_page() # Chapter 2 start
|
|
p3.insert_text((50,50), "Page 3 Content")
|
|
p4 = doc.new_page()
|
|
p4.insert_text((50,50), "Page 4 Content")
|
|
|
|
# Add TOC:
|
|
# 1. "Chap 1" -> Page 1
|
|
# 2. "Chap 2" -> Page 3
|
|
doc.set_toc([[1, "Chap 1", 1], [1, "Chap 2", 3]])
|
|
doc.save(self.pdf_path)
|
|
doc.close()
|
|
|
|
with PdfParser(self.pdf_path) as parser:
|
|
parser.process_content()
|
|
|
|
nav = parser.processed_nav_structure
|
|
|
|
# Expect 2 top level items
|
|
self.assertEqual(len(nav), 2)
|
|
self.assertEqual(nav[0]['title'], "Chap 1")
|
|
self.assertEqual(nav[0]['src'], "page_1")
|
|
self.assertEqual(nav[1]['title'], "Chap 2")
|
|
self.assertEqual(nav[1]['src'], "page_3")
|
|
|
|
# Check children of Chap 1 (Page 2 should be there)
|
|
children_c1 = nav[0]['children']
|
|
self.assertEqual(len(children_c1), 1)
|
|
# The child should likely be titled "Page 2 - Page 2 Content" or similar
|
|
self.assertIn("Page 2", children_c1[0]['title'])
|
|
self.assertEqual(children_c1[0]['src'], "page_2")
|
|
|
|
# Check children of Chap 2 (Page 4 should be there)
|
|
children_c2 = nav[1]['children']
|
|
self.assertEqual(len(children_c2), 1)
|
|
self.assertIn("Page 4", children_c2[0]['title'])
|
|
self.assertEqual(children_c2[0]['src'], "page_4")
|
|
|
|
def test_pdf_structure_without_toc(self):
|
|
# Create PDF without TOC
|
|
doc = fitz.open()
|
|
p1 = doc.new_page()
|
|
p1.insert_text((50,50), "Start")
|
|
p2 = doc.new_page()
|
|
p2.insert_text((50,50), "End")
|
|
doc.save(self.pdf_path)
|
|
doc.close()
|
|
|
|
with PdfParser(self.pdf_path) as parser:
|
|
parser.process_content()
|
|
|
|
nav = parser.processed_nav_structure
|
|
|
|
# Expect 1 top level item (Pages)
|
|
self.assertEqual(len(nav), 1)
|
|
self.assertEqual(nav[0]['title'], "Pages")
|
|
|
|
# Check children (all pages)
|
|
children = nav[0]['children']
|
|
self.assertEqual(len(children), 2)
|
|
self.assertIn("Page 1", children[0]['title'])
|
|
self.assertIn("Page 2", children[1]['title'])
|
|
|
|
def test_pdf_structure_nested_toc(self):
|
|
# Create PDF
|
|
doc = fitz.open()
|
|
p1 = doc.new_page() # Chap 1
|
|
p2 = doc.new_page() # Sec 1.1
|
|
p3 = doc.new_page() # Chap 2
|
|
|
|
doc.set_toc([
|
|
[1, "Chap 1", 1],
|
|
[2, "Sec 1.1", 2],
|
|
[1, "Chap 2", 3]
|
|
])
|
|
doc.save(self.pdf_path)
|
|
doc.close()
|
|
|
|
with PdfParser(self.pdf_path) as parser:
|
|
parser.process_content()
|
|
nav = parser.processed_nav_structure
|
|
|
|
self.assertEqual(len(nav), 2) # Chap 1, Chap 2
|
|
self.assertEqual(nav[0]['title'], "Chap 1")
|
|
|
|
# Chap 1 should have child Sec 1.1
|
|
self.assertEqual(len(nav[0]['children']), 1)
|
|
self.assertEqual(nav[0]['children'][0]['title'], "Sec 1.1")
|
|
self.assertEqual(nav[0]['children'][0]['src'], "page_2")
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|