| from __future__ import annotations |
|
|
| import base64 |
| import io |
| import unittest |
| from unittest import mock |
|
|
| from PIL import Image |
|
|
| from services import image_annotation_service |
|
|
|
|
| def _image_data_url(width: int = 200, height: int = 100) -> str: |
| output = io.BytesIO() |
| Image.new("RGB", (width, height), (245, 240, 235)).save(output, format="PNG") |
| return "data:image/png;base64," + base64.b64encode(output.getvalue()).decode("ascii") |
|
|
|
|
| class ImageAnnotationServiceTests(unittest.TestCase): |
| def test_identify_region_sends_original_and_crop_to_configured_agent(self) -> None: |
| captured: dict[str, object] = {} |
|
|
| def fake_agent(identity, system_prompt, user_prompt, image_data_urls, **kwargs): |
| captured.update( |
| identity=identity, |
| system_prompt=system_prompt, |
| user_prompt=user_prompt, |
| images=image_data_urls, |
| excluded_tokens=kwargs.get("excluded_tokens"), |
| ) |
| return ( |
| '{"label":"矿泉水瓶","object_type":"object","visible_text":"PURE",' |
| '"description":"透明塑料瓶","confidence":0.93}', |
| "vision-agent", |
| ) |
|
|
| with mock.patch.object(image_annotation_service, "call_product_visual_agent", side_effect=fake_agent): |
| result = image_annotation_service.identify_image_region( |
| {"id": "user-1"}, |
| _image_data_url(), |
| {"x": 10, "y": 20, "width": 30, "height": 40}, |
| ) |
|
|
| self.assertEqual(result["label"], "矿泉水瓶") |
| self.assertEqual(result["model"], "vision-agent") |
| self.assertEqual(result["image_width"], 200) |
| self.assertEqual(result["image_height"], 100) |
| self.assertEqual(result["bounds"], {"x": 10.0, "y": 20.0, "width": 30.0, "height": 40.0}) |
| images = captured["images"] |
| self.assertEqual(len(images), 2) |
| self.assertTrue(images[0].startswith("data:image/png;base64,")) |
| self.assertTrue(images[1].startswith("data:image/jpeg;base64,")) |
|
|
| def test_identify_region_retries_malformed_agent_output(self) -> None: |
| responses = [ |
| ("", "vision-agent"), |
| ("not json", "vision-agent"), |
| ('{"label":"人物的头发","object_type":"person_part","confidence":2}', "vision-agent"), |
| ] |
| with mock.patch.object(image_annotation_service, "call_product_visual_agent", side_effect=responses) as call: |
| result = image_annotation_service.identify_image_region( |
| {"id": "user-1"}, |
| _image_data_url(), |
| {"x": 25, "y": 10, "width": 20, "height": 25}, |
| ) |
|
|
| self.assertEqual(call.call_count, 3) |
| self.assertEqual(result["label"], "人物的头发") |
| self.assertEqual(result["confidence"], 1.0) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
|
|
