54 lines
2.0 KiB
Python
54 lines
2.0 KiB
Python
from collections import Counter
|
|
from io import BytesIO
|
|
|
|
from PIL import Image
|
|
|
|
from app import llm
|
|
|
|
|
|
def test_multi_image_ocr_retries_per_image_and_deduplicates(monkeypatch):
|
|
attempts = Counter()
|
|
|
|
def recognize(image_data, content_type):
|
|
attempts[image_data] += 1
|
|
if image_data == b"first" and attempts[image_data] == 1:
|
|
raise TimeoutError("temporary timeout")
|
|
if image_data == b"first":
|
|
return [
|
|
{"product_title": "商品A", "question": "共同问题？"},
|
|
{"product_title": "商品A", "question": "第一张的问题？"},
|
|
]
|
|
return [
|
|
{"product_title": "商品A", "question": "共同问题？"},
|
|
{"product_title": "商品A", "question": "第二张的问题？"},
|
|
]
|
|
|
|
monkeypatch.setattr(llm.settings, "llm_api_key", "test-key")
|
|
monkeypatch.setattr(llm, "_compress_ocr_image", lambda data, content_type: (data, content_type))
|
|
monkeypatch.setattr(llm, "_extract_questions_from_image", recognize)
|
|
monkeypatch.setattr(llm.time, "sleep", lambda seconds: None)
|
|
|
|
result = llm.extract_questions_from_images([
|
|
(b"first", "image/png"),
|
|
(b"second", "image/png"),
|
|
])
|
|
|
|
assert attempts == Counter({b"first": 2, b"second": 1})
|
|
assert [item["question"] for item in result] == ["共同问题？", "第一张的问题？", "第二张的问题？"]
|
|
assert [item["source_index"] for item in result] == [0, 0, 1]
|
|
|
|
|
|
def test_ocr_image_compression_limits_dimensions_and_converts_to_jpeg():
|
|
image = Image.new("RGBA", (3200, 2600), (20, 100, 220, 180))
|
|
source = BytesIO()
|
|
image.save(source, format="PNG")
|
|
|
|
compressed, content_type = llm._compress_ocr_image(source.getvalue(), "image/png")
|
|
|
|
with Image.open(BytesIO(compressed)) as result:
|
|
assert max(result.size) <= 2200
|
|
assert result.mode == "RGB"
|
|
assert result.format == "JPEG"
|
|
assert content_type == "image/jpeg"
|
|
assert len(compressed) < len(source.getvalue())
|