187 lines
9.9 KiB
Python
187 lines
9.9 KiB
Python
import json
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pymupdf
|
|
from PIL import Image
|
|
|
|
from copienator import EvaluationWorkspace, ExitCode, atomic_write_json, read_json
|
|
from copienator.annotation_layout import fingerprint
|
|
from copienator.commands import gemini_for_labels as labels, repair_orientation as repair, cutleft
|
|
from copienator.commands.page_splitter import PDFPreviewer
|
|
from copienator.orientation import (
|
|
incoherent, load_provenance, record_detection, repair_plan, report_path,
|
|
reports, resolve_unknowns, save_provenance,
|
|
)
|
|
|
|
|
|
class OrientationTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.tmp.cleanup)
|
|
self.workspace = EvaluationWorkspace(Path(self.tmp.name))
|
|
self.workspace.copies_dir.mkdir()
|
|
self.workspace.cutleft_dir.mkdir()
|
|
self.pdf = self.workspace.copies_dir / 'Copie01.pdf'
|
|
with pymupdf.open() as doc:
|
|
for i in range(4):
|
|
page = doc.new_page(width=200, height=300)
|
|
page.insert_text((30, 30), f'Page {i + 1}')
|
|
doc.save(self.pdf)
|
|
self.origins = [{'origin': p} for p in (1, 2, 2, 1)]
|
|
save_provenance(self.workspace, self.pdf, self.origins)
|
|
atomic_write_json(self.workspace.cutleft_dir / 'Copie01_schema.json', {
|
|
'columns_per_file': [2, 2], 'source_sha256': fingerprint(self.pdf),
|
|
})
|
|
self.images = [self.workspace.cutleft_dir / f'Copie01_{i:02d}.jpg' for i in (1, 2)]
|
|
for image in self.images:
|
|
Image.new('RGB', (400, 300), 'white').save(image)
|
|
|
|
def flag(self, pages):
|
|
data = {'sha256': fingerprint(self.pdf), 'images': {'test.jpg': pages}}
|
|
atomic_write_json(report_path(self.workspace, self.pdf.stem), data)
|
|
return data
|
|
|
|
def test_majority_excludes_unknowns_and_ties_remain_unknown(self):
|
|
for votes, expected in (
|
|
(['upside_down', 'unknown', 'unknown'], ['upside_down'] * 3),
|
|
(['upright', 'unknown'], ['upright'] * 2),
|
|
(['upright', 'upside_down', 'unknown'], ['upright', 'upside_down', 'unknown']),
|
|
(['unknown', 'unknown'], ['unknown', 'unknown']),
|
|
):
|
|
with self.subTest(votes=votes):
|
|
verdict = dict(enumerate(votes, 1))
|
|
actual = resolve_unknowns(verdict, [list(verdict)])
|
|
self.assertEqual(list(actual.values()), expected)
|
|
self.assertEqual(incoherent(actual, list(actual)), 'upright' in expected and 'upside_down' in expected)
|
|
|
|
def test_plan_reverses_only_coherent_origin_fragments_in_place(self):
|
|
verdict = {1: 'upside_down', 2: 'upright', 3: 'upright', 4: 'upside_down'}
|
|
self.assertEqual(repair_plan(verdict, [[1, 4], [2, 3]]), ([4, 2, 3, 1], {1, 4}))
|
|
verdict[4] = 'upright'
|
|
self.assertEqual(repair_plan(verdict, [[1, 4], [2, 3]]), ([1, 2, 3, 4], {1}))
|
|
|
|
def test_montage_local_columns_map_to_copy_pages_and_reject_bad_indices(self):
|
|
record_detection(self.workspace, self.images[1], [2])
|
|
data = reports(self.workspace)['Copie01.pdf']
|
|
self.assertEqual(data['images'][self.images[1].name], [4])
|
|
with self.assertRaises(ValueError):
|
|
record_detection(self.workspace, self.images[1], [3])
|
|
record_detection(self.workspace, self.images[1], [])
|
|
self.assertFalse(reports(self.workspace))
|
|
|
|
def test_orientation_report_takes_priority_over_crop_error(self):
|
|
client = Mock()
|
|
client.models.generate_content.side_effect = [
|
|
Mock(text=json.dumps({'upside_down_pages': [2], 'error': 'insufficient_crop'})),
|
|
Mock(text=json.dumps({'name': 'Continued', 'list': []})),
|
|
]
|
|
with patch.object(cutleft, 'regenerate_full_page_image') as fallback:
|
|
labels.process_copy_group(self.workspace, 'Copie01', self.images,
|
|
client=client, labels_text='Ex 1', names_text='Student', valid_labels={'Ex 1'},
|
|
valid_names={'Unknown', 'Continued'}, overwrite=True, sleep=lambda _: None,
|
|
target_interval=0)
|
|
fallback.assert_not_called()
|
|
self.assertFalse((self.workspace.copies_dir / 'Copie01_01.json').exists())
|
|
self.assertTrue((self.workspace.copies_dir / 'Copie01_02.json').exists())
|
|
|
|
def test_coherent_repair_avoids_second_request_and_preserves_content(self):
|
|
data = self.flag([1, 4])
|
|
with patch.object(repair, 'verify_pages') as verify:
|
|
repair.prepare_repair(self.workspace, self.pdf, data, Mock())
|
|
verify.assert_not_called()
|
|
self.assertEqual(data['order'], [4, 2, 3, 1])
|
|
self.assertFalse(data['warning'])
|
|
with pymupdf.open(report_path(self.workspace, 'Copie01').with_suffix('.prepared.pdf')) as doc:
|
|
self.assertEqual([p.get_text().strip() for p in doc], ['Page 4', 'Page 2', 'Page 3', 'Page 1'])
|
|
self.assertGreater(doc[0].get_text('blocks')[0][1], 200)
|
|
self.assertLess(doc[1].get_text('blocks')[0][1], 50)
|
|
|
|
def test_full_page_verdict_is_definitive_and_unknown_inherits_origin(self):
|
|
data = self.flag([1])
|
|
with patch.object(repair, 'verify_pages', return_value={1: 'upside_down', 4: 'unknown'}) as verify:
|
|
repair.prepare_repair(self.workspace, self.pdf, data, Mock())
|
|
self.assertEqual(verify.call_args.args[2], [1, 4])
|
|
self.assertEqual(data['rotated'], [1, 4])
|
|
self.assertFalse(data['warning'])
|
|
|
|
def test_only_final_mixed_verdict_warns_about_order(self):
|
|
for final, mixed, rotate in (
|
|
({1: 'upright', 4: 'upright'}, False, []),
|
|
({1: 'upright', 4: 'upside_down'}, True, [4]),
|
|
({1: 'unknown', 4: 'unknown'}, False, []),
|
|
):
|
|
with self.subTest(final=final):
|
|
data = self.flag([1])
|
|
prepared = report_path(self.workspace, 'Copie01').with_suffix('.prepared.pdf')
|
|
prepared.unlink(missing_ok=True)
|
|
with patch.object(repair, 'verify_pages', return_value=final):
|
|
repair.prepare_repair(self.workspace, self.pdf, data, Mock())
|
|
self.assertEqual('ordre des pages' in data['warning'], mixed)
|
|
self.assertEqual(data['rotated'], rotate)
|
|
|
|
def test_failed_verification_keeps_disputed_group_but_fixes_coherent_group(self):
|
|
data = self.flag([1, 2, 3])
|
|
with patch.object(repair, 'verify_pages', side_effect=ValueError('bad JSON')):
|
|
repair.prepare_repair(self.workspace, self.pdf, data, Mock())
|
|
self.assertEqual(data['rotated'], [2, 3])
|
|
self.assertNotIn('ordre des pages', data['warning'])
|
|
self.assertIn('indéterminée', data['warning'])
|
|
|
|
def test_full_page_verification_validates_ids_and_uses_low_thinking(self):
|
|
client = Mock()
|
|
client.models.generate_content.return_value.text = json.dumps({'pages': [
|
|
{'page': 1, 'orientation': 'upright'}, {'page': 4, 'orientation': 'unknown'}]})
|
|
self.assertEqual(repair.verify_pages(client, self.pdf, [1, 4]), {1: 'upright', 4: 'unknown'})
|
|
request = client.models.generate_content.call_args.kwargs
|
|
self.assertEqual(request['config'].thinking_config.thinking_level.value.lower(), 'low')
|
|
self.assertEqual(sum(p.inline_data is not None for p in request['contents'][0].parts), 2)
|
|
client.models.generate_content.return_value.text = '{"pages": [{"page": 9, "orientation": "upright"}]}'
|
|
with self.assertRaises(ValueError):
|
|
repair.verify_pages(client, self.pdf, [1, 4])
|
|
|
|
def test_resume_does_not_rotate_twice_and_later_alert_is_manual(self):
|
|
data = self.flag([1, 4])
|
|
def finish_labels(workspace, targets, **kwargs):
|
|
atomic_write_json(workspace.copies_dir / 'Copie01_01.json', {'name': 'Student', 'list': []})
|
|
return ExitCode.SUCCESS
|
|
result = (Image.new('RGB', (100, 100)), [Image.new('RGB', (100, 100))], {'columns_per_file': [4]})
|
|
with patch.object(cutleft, 'process_single_pdf', side_effect=RuntimeError('interrupted')):
|
|
with self.assertRaises(RuntimeError):
|
|
repair.repair_copy(self.workspace, self.pdf, data, Mock())
|
|
rotated_digest = fingerprint(self.pdf)
|
|
journal = reports(self.workspace)['Copie01.pdf']
|
|
with patch.object(cutleft, 'process_single_pdf', return_value=result), patch.object(labels, 'run', side_effect=finish_labels):
|
|
self.assertEqual(repair.repair_copy(self.workspace, self.pdf, journal, Mock()), ExitCode.SUCCESS)
|
|
self.assertEqual(fingerprint(self.pdf), rotated_digest)
|
|
self.assertFalse(reports(self.workspace))
|
|
record_detection(self.workspace, self.images[0], [1])
|
|
journal = reports(self.workspace)['Copie01.pdf']
|
|
with patch.object(repair, 'prepare_repair') as prepare:
|
|
self.assertEqual(repair.repair_copy(self.workspace, self.pdf, journal, Mock()), ExitCode.PARTIAL)
|
|
prepare.assert_not_called()
|
|
|
|
def test_splitter_preserves_original_ids_through_reversal_and_secondary_splits(self):
|
|
preview = PDFPreviewer.__new__(PDFPreviewer)
|
|
preview.doc = pymupdf.open(self.pdf)
|
|
preview.original_page_ids = {p.xref: i + 1 for i, p in enumerate(preview.doc)}
|
|
preview.doc.select([3, 2, 1, 0])
|
|
preview.file_rotation = preview.global_rotation = 0
|
|
preview.base_name = 'Copie01'
|
|
preview.split_dir = self.workspace.root / 'split'
|
|
preview.reorder_dir = self.workspace.root / 'reorder'
|
|
preview.final_file = self.workspace.root / 'result.pdf'
|
|
preview.output_dir = None
|
|
preview.page_settings = [{'keep': 'as_is', 'line_x': 100, 'rotation': 0,
|
|
'secondary_splitters': [75]} for _ in range(4)]
|
|
preview.split_pdf()
|
|
preview.reorder_pdfs()
|
|
preview.concate_files()
|
|
self.assertEqual([p['origin'] for p in preview.output_origins], [4, 4, 3, 3, 2, 2, 1, 1])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|