import json import tempfile import unittest from pathlib import Path from unittest.mock import Mock, patch from PIL import Image from copienator import EvaluationWorkspace, ExitCode, read_json from copienator import name_grids as names from copienator.commands import cutleft, gemini_for_labels as labels class NameGridTests(unittest.TestCase): def setUp(self): temporary = tempfile.TemporaryDirectory() self.addCleanup(temporary.cleanup) self.workspace = EvaluationWorkspace(Path(temporary.name)) self.workspace.copies_dir.mkdir() self.workspace.cutleft_dir.mkdir() self.directory = self.workspace.cutleft_dir / 'Names' (self.workspace.root / 'names').write_text('Alice\nBob\n') (self.workspace.root / 'labels').write_text('Ex 1\n') self.client = Mock() def crop(self, number): pdf = self.workspace.copies_dir / f'Copie{number:02d}.pdf' pdf.write_bytes(b'pdf') page = Image.new('RGB', (1654, 2338), 'white') page.putpixel((0, 0), (0, 0, 0)) names.save_crop(page, pdf, self.directory, 200) (self.workspace.cutleft_dir / f'{pdf.stem}_01.jpg').write_bytes(b'image') return pdf def identify(self, groups, overwrite=False): return names.identify_names(self.workspace, groups, client=self.client, model='test', names_text='Alice\nBob\n', overwrite=overwrite) def test_crop_geometry_and_partial_grid_mapping(self): for number in range(1, 26): self.crop(number) with Image.open(self.directory / 'Copie01.jpg') as crop: self.assertEqual(crop.size, (220, 584)) self.assertLess(crop.getpixel((0, 0))[0], 80) grids = names.build_grids(self.directory, self.workspace.copies_dir) self.assertEqual([len(grid['cells']) for grid in grids], [24, 1]) self.assertEqual(grids[0]['cells'][6], {'cell': 7, 'row': 2, 'column': 1, 'copy': 'Copie07'}) self.assertEqual(grids[1]['cells'][0]['copy'], 'Copie25') with Image.open(self.directory / 'Grids' / grids[1]['image']) as grid: self.assertEqual(grid.size, (1362, 2494)) self.assertLess(grid.getpixel((228, 500))[0], 40) # Deleted/stale PDFs cannot be assigned names from old crops. (self.workspace.copies_dir / 'Copie25.pdf').unlink() (self.workspace.copies_dir / 'Copie24.pdf').write_bytes(b'changed') grids = names.build_grids(self.directory, self.workspace.copies_dir) self.assertEqual(len(grids), 1) self.assertEqual(len(grids[0]['cells']), 23) self.assertFalse((self.directory / 'Grids' / 'grid_002.jpg').exists()) def test_validation_cache_and_roster_changes(self): for number in range(1, 5): self.crop(number) self.client.models.generate_content.return_value = Mock(text=json.dumps({'cells': [ {'cell': 1, 'name': 'Alice'}, {'cell': 2, 'name': 'Outsider'}, {'cell': 3, 'name': 'Bob'}, {'cell': 3, 'name': 'Alice'}, {'cell': 24, 'name': 'Bob'}, ]})) groups = {f'Copie{i:02d}': [] for i in range(1, 5)} expected = {'Copie01': 'Alice', **{f'Copie{i:02d}': 'Missing' for i in range(2, 5)}} self.assertEqual(self.identify(groups), expected) request = self.client.models.generate_content.call_args.kwargs self.assertEqual(request['config'].thinking_config.thinking_level.value, 'LOW') text = request['contents'][0].parts[1].text self.assertIn('6 columns and 4 rows', text) self.assertIn('Cell 4: row 1, column 4', text) self.assertNotIn('Cell 5:', text) self.assertEqual(self.identify(groups), expected) self.assertEqual(self.client.models.generate_content.call_count, 1) self.identify(groups, overwrite=True) self.assertEqual(self.client.models.generate_content.call_count, 2) names.identify_names(self.workspace, groups, client=self.client, model='test', names_text='Bob\n', overwrite=False) self.assertEqual(read_json(self.directory / 'Copie01_result.json')['name'], 'Missing') def test_api_failure_is_retried_on_resume(self): self.crop(1) self.client.models.generate_content.side_effect = TimeoutError() self.assertEqual(self.identify({'Copie01': []}), {'Copie01': 'Missing'}) self.assertFalse((self.directory / 'Copie01_result.json').exists()) self.identify({'Copie01': []}) self.assertEqual(self.client.models.generate_content.call_count, 2) def test_known_and_missing_names_choose_different_label_requests(self): for number in (1, 2): self.crop(number) (self.workspace.cutleft_dir / 'Copie01_02.jpg').write_bytes(b'image') self.client.models.generate_content.side_effect = [ Mock(text='{"cells":[{"cell":1,"name":"Alice"},{"cell":2,"name":"Missing"}]}'), Mock(text='{"list":[{"box_2d":[1,2,3,4],"label":"Ex 1"}]}'), Mock(text='{"list":[]}'), Mock(text='{"name":"Bob","list":[]}'), ] with patch.object(labels.config, 'MODEL_FOR_NAME_ID', 'custom-name-model'), \ patch.object(labels.config, 'THINKING_LEVEL_FOR_NAME', 'medium'): status = labels.run(self.workspace, [self.workspace.root], client=self.client, sleep=lambda _: None, max_workers=1) self.assertEqual(status, ExitCode.SUCCESS) requests = self.client.models.generate_content.call_args_list self.assertEqual(len(requests), 4) self.assertEqual(requests[0].kwargs['model'], 'custom-name-model') self.assertEqual(requests[0].kwargs['config'].thinking_config.thinking_level.value, 'MEDIUM') self.assertEqual(requests[1].kwargs['model'], labels.MODEL_ID) for request in requests[1:3]: text = request.kwargs['contents'][0].parts[1].text self.assertNotIn('Alice', text) self.assertNotIn('Bob', text) self.assertNotIn('student name', text) self.assertNotIn('name', request.kwargs['config'].response_json_schema['properties']) self.assertIn('Alice', requests[3].kwargs['contents'][0].parts[1].text) for stem, expected in [('Copie01_01', 'Alice'), ('Copie01_02', 'Continued'), ('Copie02_01', 'Bob')]: self.assertEqual(read_json(self.workspace.copies_dir / f'{stem}.json')['name'], expected) self.client.models.generate_content.reset_mock() labels.run(self.workspace, [self.workspace.root], client=self.client, sleep=lambda _: None) self.client.models.generate_content.assert_not_called() def test_save_label_crops_also_saves_name_grid(self): pdf = self.workspace.copies_dir / 'Copie01.pdf' pdf.write_bytes(b'pdf') page = Image.new('RGB', (1654, 2338), 'white') with patch.object(cutleft, 'get_pdf_pages', return_value=[page]): result = cutleft.process_single_pdf(pdf, automatic_margin=False) cutleft.save_results(result, pdf, self.workspace.cutleft_dir) self.assertTrue((self.directory / 'Copie01.jpg').is_file()) self.assertEqual(read_json(self.directory / 'Grids' / 'manifest.json')[0]['cells'][0]['copy'], 'Copie01') def test_known_name_survives_label_retry(self): self.crop(1) image = self.workspace.cutleft_dir / 'Copie01_01.jpg' self.client.models.generate_content.side_effect = [TimeoutError(), Mock(text='{"list":[]}')] labels.process_copy_group(self.workspace, 'Copie01', [image], client=self.client, labels_text='Ex 1', names_text='Alice\nBob', valid_labels={'Ex 1'}, valid_names={'Alice', 'Bob'}, overwrite=False, known_name='Alice', sleep=lambda _: None, target_interval=0) self.assertEqual(read_json(self.workspace.copies_dir / 'Copie01_01.json')['name'], 'Alice') for call in self.client.models.generate_content.call_args_list: self.assertNotIn('Alice', call.kwargs['contents'][0].parts[1].text)