latest
This commit is contained in:
@@ -0,0 +1,210 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Simple test script for split_map.py functionality.
|
||||
Run this to verify the split map generator works correctly.
|
||||
"""
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
# Mock implementations for testing without full environment
|
||||
class MockStorage:
|
||||
def __init__(self):
|
||||
self.files = {}
|
||||
|
||||
def download_file(self, bucket, path):
|
||||
return self.files.get(f"{bucket}/{path}", b"{}")
|
||||
|
||||
def upload_file(self, bucket, path, data, mime, upsert=True):
|
||||
self.files[f"{bucket}/{path}"] = data
|
||||
|
||||
class MockClient:
|
||||
def __init__(self):
|
||||
self.files = []
|
||||
self.artefacts = []
|
||||
|
||||
@property
|
||||
def supabase(self):
|
||||
return self
|
||||
|
||||
def table(self, name):
|
||||
self.current_table = name
|
||||
return self
|
||||
|
||||
def select(self, cols):
|
||||
return self
|
||||
|
||||
def eq(self, col, val):
|
||||
return self
|
||||
|
||||
def single(self):
|
||||
return self
|
||||
|
||||
def order(self, col, desc=False):
|
||||
return self
|
||||
|
||||
def execute(self):
|
||||
if self.current_table == 'files':
|
||||
return type('Result', (), {'data': {
|
||||
'id': 'test-file-id',
|
||||
'bucket': 'test-bucket',
|
||||
'cabinet_id': 'test-cabinet',
|
||||
'name': 'test.pdf',
|
||||
'path': 'test/path.pdf',
|
||||
'mime_type': 'application/pdf'
|
||||
}})()
|
||||
elif self.current_table == 'document_artefacts':
|
||||
return type('Result', (), {'data': self.artefacts})()
|
||||
return type('Result', (), {'data': []})()
|
||||
|
||||
def insert(self, data):
|
||||
if self.current_table == 'document_artefacts':
|
||||
self.artefacts.append(data)
|
||||
return self
|
||||
|
||||
def test_outline_extraction():
|
||||
"""Test PDF outline extraction (requires PyMuPDF)."""
|
||||
print("Testing PDF outline extraction...")
|
||||
|
||||
try:
|
||||
from routers.database.files.split_map import _try_outline
|
||||
|
||||
# Create a simple PDF with bookmarks (this would need actual PDF bytes in real usage)
|
||||
# For now, just test the import and function signature
|
||||
result = _try_outline(b"dummy pdf bytes")
|
||||
print(f"Outline extraction result: {result}")
|
||||
|
||||
except ImportError:
|
||||
print("PyMuPDF not available - outline extraction will be skipped")
|
||||
except Exception as e:
|
||||
print(f"Outline extraction test failed: {e}")
|
||||
|
||||
def test_heading_extraction():
|
||||
"""Test heading extraction from Docling JSON."""
|
||||
print("\nTesting heading extraction...")
|
||||
|
||||
from routers.database.files.split_map import _try_headings
|
||||
|
||||
# Mock Docling JSON with headings
|
||||
docling_json = {
|
||||
"blocks": [
|
||||
{
|
||||
"role": "heading",
|
||||
"text": "Chapter 1: Introduction",
|
||||
"page": 1,
|
||||
"type": "h1"
|
||||
},
|
||||
{
|
||||
"role": "paragraph",
|
||||
"text": "Some content...",
|
||||
"page": 1
|
||||
},
|
||||
{
|
||||
"role": "heading",
|
||||
"text": "Chapter 2: Methods",
|
||||
"page": 15,
|
||||
"type": "h1"
|
||||
},
|
||||
{
|
||||
"role": "heading",
|
||||
"text": "Chapter 3: Results",
|
||||
"page": 30,
|
||||
"type": "h1"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
result = _try_headings(docling_json)
|
||||
print(f"Headings extraction result: {result}")
|
||||
|
||||
expected = [("Chapter 1: Introduction", 1, 1), ("Chapter 2: Methods", 15, 1), ("Chapter 3: Results", 30, 1)]
|
||||
if result == expected:
|
||||
print("✓ Heading extraction test passed")
|
||||
else:
|
||||
print(f"✗ Heading extraction test failed. Expected: {expected}, Got: {result}")
|
||||
|
||||
def test_toc_extraction():
|
||||
"""Test TOC extraction from Tika text."""
|
||||
print("\nTesting TOC extraction...")
|
||||
|
||||
from routers.database.files.split_map import _try_toc_text
|
||||
|
||||
# Mock Tika text with TOC
|
||||
tika_text = """
|
||||
Table of Contents
|
||||
|
||||
Introduction .................. 1
|
||||
Chapter 1: Getting Started .... 5
|
||||
Chapter 2: Advanced Topics .... 15
|
||||
Chapter 3: Examples ........... 25
|
||||
Chapter 4: Conclusion ......... 35
|
||||
Appendix A .................... 40
|
||||
Index ......................... 45
|
||||
|
||||
Some other content follows...
|
||||
"""
|
||||
|
||||
result = _try_toc_text(tika_text)
|
||||
print(f"TOC extraction result: {result}")
|
||||
|
||||
if result and len(result) >= 5:
|
||||
print("✓ TOC extraction test passed")
|
||||
else:
|
||||
print("✗ TOC extraction test failed")
|
||||
|
||||
def test_entries_building():
|
||||
"""Test building entries from start points."""
|
||||
print("\nTesting entries building...")
|
||||
|
||||
from routers.database.files.split_map import _entries_from_starts, _entries_from_pairs
|
||||
|
||||
# Test starts format
|
||||
starts = [("Chapter 1", 1, 1), ("Chapter 2", 15, 1), ("Chapter 3", 30, 1)]
|
||||
entries1 = _entries_from_starts(starts, 50, "headings")
|
||||
print(f"Entries from starts: {len(entries1)} entries")
|
||||
|
||||
# Test pairs format
|
||||
pairs = [("Introduction", 1), ("Methods", 15), ("Results", 30)]
|
||||
entries2 = _entries_from_pairs(pairs, 50, "outline")
|
||||
print(f"Entries from pairs: {len(entries2)} entries")
|
||||
|
||||
# Verify structure
|
||||
if entries1 and all(key in entries1[0] for key in ['id', 'title', 'start_page', 'end_page', 'confidence']):
|
||||
print("✓ Entry building test passed")
|
||||
else:
|
||||
print("✗ Entry building test failed")
|
||||
|
||||
def test_tiny_section_merging():
|
||||
"""Test merging of tiny sections."""
|
||||
print("\nTesting tiny section merging...")
|
||||
|
||||
from routers.database.files.split_map import _entries_from_pairs
|
||||
|
||||
# Create pairs with tiny sections (1-2 pages)
|
||||
pairs = [("Chapter 1", 1), ("Short Section", 2), ("Another Short", 3), ("Chapter 2", 15)]
|
||||
entries = _entries_from_pairs(pairs, 50, "test")
|
||||
|
||||
print(f"Original pairs: {len(pairs)}, Final entries: {len(entries)}")
|
||||
|
||||
# Should merge tiny sections
|
||||
if len(entries) < len(pairs):
|
||||
print("✓ Tiny section merging test passed")
|
||||
else:
|
||||
print("✗ Tiny section merging test failed")
|
||||
|
||||
def main():
|
||||
"""Run all tests."""
|
||||
print("Running split_map.py tests...\n")
|
||||
|
||||
test_outline_extraction()
|
||||
test_heading_extraction()
|
||||
test_toc_extraction()
|
||||
test_entries_building()
|
||||
test_tiny_section_merging()
|
||||
|
||||
print("\nTests completed!")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user