146 lines
8.0 KiB
Python
146 lines
8.0 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, 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)
|