46 lines
2.1 KiB
Python
46 lines
2.1 KiB
Python
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from jsonschema import Draft202012Validator
|
|
from pydantic import ValidationError
|
|
|
|
from copienator.commands import gemini_for_labels as labels
|
|
|
|
|
|
class LabelSchemaTests(unittest.TestCase):
|
|
def test_request_schema_and_parser_agree_on_success_and_recovery(self):
|
|
cases = [
|
|
({}, False),
|
|
({'name': 'Student'}, False),
|
|
({'name': 'Student', 'error': None, 'upside_down_pages': []}, False),
|
|
({'name': 'Student', 'list': []}, True),
|
|
({'name': 'Student', 'list': [
|
|
{'box_2d': [1, 2, 3, 4], 'label': 'Ex 1'}]}, True),
|
|
({'error': 'insufficient_crop'}, True),
|
|
({'upside_down_pages': [1]}, True),
|
|
({'error': 'insufficient_crop', 'upside_down_pages': [1]}, True),
|
|
({'error': None}, False),
|
|
({'upside_down_pages': []}, False),
|
|
({'upside_down_pages': [0]}, False),
|
|
({'error': 'other'}, False),
|
|
({'name': 'Student', 'list': None}, False),
|
|
]
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
image = Path(directory) / 'Copie01_01.jpg'
|
|
image.write_bytes(b'image')
|
|
for known_name, model in ((None, labels.AnnotationData), ('Student', labels.LabelData)):
|
|
_, config = labels.generate_request(
|
|
image, 'Ex 1', 'Student', [], [], known_name=known_name)
|
|
schema = config.response_json_schema
|
|
Draft202012Validator.check_schema(schema)
|
|
validator = Draft202012Validator(schema)
|
|
for payload, expected in cases + [({'list': []}, bool(known_name))]:
|
|
with self.subTest(known_name=known_name, payload=payload):
|
|
self.assertEqual(validator.is_valid(payload), expected)
|
|
if expected:
|
|
model.model_validate(payload)
|
|
else:
|
|
with self.assertRaises(ValidationError):
|
|
model.model_validate(payload)
|