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()