108 lines
5.9 KiB
Python
108 lines
5.9 KiB
Python
import json
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import Mock, patch
|
|
|
|
from PIL import Image
|
|
|
|
from copienator import EvaluationWorkspace, read_json
|
|
from copienator.commands import cutleft, gemini_for_labels as labels
|
|
|
|
|
|
class FullPageFallbackTests(unittest.TestCase):
|
|
def setUp(self):
|
|
temporary = tempfile.TemporaryDirectory()
|
|
self.addCleanup(temporary.cleanup)
|
|
self.workspace = EvaluationWorkspace(Path(temporary.name))
|
|
self.workspace.cutleft_dir.mkdir()
|
|
self.workspace.copies_dir.mkdir()
|
|
self.images = [self.workspace.cutleft_dir / f'Copie01_{i:02d}.jpg' for i in (1, 2, 3)]
|
|
for image in self.images:
|
|
image.write_bytes(b'original crop')
|
|
self.schema_path = self.workspace.cutleft_dir / 'Copie01_schema.json'
|
|
self.schema_path.write_text(json.dumps({'columns_per_file': [1, 2, 1],
|
|
'full_width_pages': [], 'total_pages': 4}))
|
|
self.pages = [Image.new('RGB', (100, 200), color) for color in ('red', 'green', 'blue', 'white')]
|
|
self.client = Mock()
|
|
self.error = Mock(text='{"error":"insufficient_crop","error_reason":"Exercise heading clipped"}')
|
|
self.success = Mock(text='{"name":"Student","list":[{"box_2d":[1,2,3,4],"label":"Ex 1"}]}')
|
|
|
|
def process(self):
|
|
return labels.process_copy_group(
|
|
self.workspace, 'Copie01', self.images, client=self.client,
|
|
labels_text='Ex 1', names_text='Student', valid_labels={'Ex 1'},
|
|
valid_names={'Student'}, overwrite=False, sleep=lambda _: None, target_interval=0,
|
|
)
|
|
|
|
def test_regenerates_only_affected_pages_and_updates_metadata(self):
|
|
with patch.object(cutleft, 'get_pdf_pages', return_value=self.pages):
|
|
cutleft.regenerate_full_page_image(self.workspace, self.images[1])
|
|
self.assertEqual(self.images[0].read_bytes(), b'original crop')
|
|
self.assertEqual(self.images[2].read_bytes(), b'original crop')
|
|
with Image.open(self.images[1]) as result:
|
|
self.assertEqual(result.size, (205, 200))
|
|
self.assertGreater(result.getpixel((40, 50))[1], 100)
|
|
self.assertGreater(result.getpixel((150, 50))[2], 200)
|
|
schema = read_json(self.schema_path)
|
|
self.assertEqual(schema['columns_per_file'], [1, 2, 1])
|
|
self.assertEqual(schema['full_width_pages'], [2, 3])
|
|
self.assertEqual([crop['right_px'] for crop in schema['page_crops']], [100, 100])
|
|
with self.assertRaisesRegex(ValueError, 'already contains full pages'):
|
|
cutleft.regenerate_full_page_image(self.workspace, self.images[1])
|
|
|
|
def test_error_regenerates_and_resends_with_context(self):
|
|
self.client.models.generate_content.side_effect = [self.success, self.error, self.success, self.success]
|
|
with patch.object(cutleft, 'get_pdf_pages', return_value=self.pages):
|
|
self.assertEqual(self.process(), 3)
|
|
calls = self.client.models.generate_content.call_args_list
|
|
retry_parts = calls[2].kwargs['contents'][0].parts
|
|
self.assertEqual(retry_parts[0].inline_data.data, self.images[1].read_bytes())
|
|
self.assertIn('now contains full pages', retry_parts[1].text)
|
|
self.assertIn('"Ex 1"', retry_parts[1].text)
|
|
result = read_json(self.workspace.copies_dir / 'Copie01_02.json')
|
|
self.assertEqual(set(result), {'name', 'list'})
|
|
self.assertEqual(result['list'][0]['label'], 'Ex 1')
|
|
|
|
def test_second_crop_error_stops_and_defers_following_images(self):
|
|
self.client.models.generate_content.side_effect = [self.error, self.error]
|
|
with patch.object(cutleft, 'get_pdf_pages', return_value=self.pages), \
|
|
self.assertRaises(labels.IncompleteCopyError) as caught:
|
|
self.process()
|
|
self.assertEqual(caught.exception.images, self.images)
|
|
self.assertEqual(self.client.models.generate_content.call_count, 2)
|
|
self.assertFalse(list(self.workspace.copies_dir.glob('*.json')))
|
|
|
|
def test_bad_mapping_leaves_original_image_and_schema_intact(self):
|
|
original_schema = self.schema_path.read_bytes()
|
|
self.client.models.generate_content.return_value = self.error
|
|
with patch.object(cutleft, 'get_pdf_pages', return_value=self.pages[:1]), \
|
|
self.assertRaises(labels.IncompleteCopyError):
|
|
self.process()
|
|
self.assertEqual(self.client.models.generate_content.call_count, 1)
|
|
self.assertEqual(self.schema_path.read_bytes(), original_schema)
|
|
self.assertEqual(self.images[0].read_bytes(), b'original crop')
|
|
|
|
def test_full_page_retry_gets_fresh_bounded_attempt_budget(self):
|
|
self.client.models.generate_content.side_effect = [TimeoutError(), TimeoutError(), self.error,
|
|
TimeoutError(), TimeoutError(), self.success,
|
|
self.success, self.success]
|
|
with patch.object(cutleft, 'get_pdf_pages', return_value=self.pages):
|
|
self.assertEqual(self.process(), 3)
|
|
self.assertEqual(self.client.models.generate_content.call_count, 8)
|
|
|
|
def test_resumed_full_pages_use_correct_prompt_without_regenerating(self):
|
|
schema = read_json(self.schema_path)
|
|
schema['full_width_pages'] = [1]
|
|
self.schema_path.write_text(json.dumps(schema))
|
|
self.client.models.generate_content.return_value = self.error
|
|
with patch.object(cutleft, 'get_pdf_pages') as render, self.assertRaises(labels.IncompleteCopyError):
|
|
self.process()
|
|
render.assert_not_called()
|
|
prompt = self.client.models.generate_content.call_args.kwargs['contents'][0].parts[1].text
|
|
self.assertIn('now contains full pages', prompt)
|
|
|
|
def test_missing_success_payload_is_not_accepted(self):
|
|
with self.assertRaises(ValueError):
|
|
labels.AnnotationData.model_validate_json('{}')
|