103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
|
import io
|
||
|
|
import json
|
||
|
|
import pytest
|
||
|
|
import numpy as np
|
||
|
|
import cv2
|
||
|
|
from unittest.mock import AsyncMock
|
||
|
|
|
||
|
|
from app.core.exceptions import LLMParseError
|
||
|
|
from app.models.image_models import ImageExtractRequest
|
||
|
|
|
||
|
|
|
||
|
|
def _make_test_image_bytes(width=100, height=80) -> bytes:
|
||
|
|
img = np.zeros((height, width, 3), dtype=np.uint8)
|
||
|
|
img[10:50, 10:60] = (255, 0, 0) # blue rectangle
|
||
|
|
_, buf = cv2.imencode(".jpg", img)
|
||
|
|
return buf.tobytes()
|
||
|
|
|
||
|
|
|
||
|
|
SAMPLE_QUADS_JSON = json.dumps([
|
||
|
|
{
|
||
|
|
"subject": "电缆接头",
|
||
|
|
"predicate": "位于",
|
||
|
|
"object": "配电箱左侧",
|
||
|
|
"qualifier": "2024年检修",
|
||
|
|
"bbox": {"x": 10, "y": 10, "w": 40, "h": 30},
|
||
|
|
}
|
||
|
|
])
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def image_bytes():
|
||
|
|
return _make_test_image_bytes()
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def req():
|
||
|
|
return ImageExtractRequest(file_path="image/test.jpg", task_id=1)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_extract_quads_returns_items(mock_llm, mock_storage, image_bytes, req):
|
||
|
|
mock_storage.download_bytes = AsyncMock(return_value=image_bytes)
|
||
|
|
mock_llm.chat_vision = AsyncMock(return_value=SAMPLE_QUADS_JSON)
|
||
|
|
mock_storage.upload_bytes = AsyncMock(return_value=None)
|
||
|
|
|
||
|
|
from app.services.image_service import extract_quads
|
||
|
|
result = await extract_quads(req, mock_llm, mock_storage)
|
||
|
|
|
||
|
|
assert len(result.items) == 1
|
||
|
|
item = result.items[0]
|
||
|
|
assert item.subject == "电缆接头"
|
||
|
|
assert item.predicate == "位于"
|
||
|
|
assert item.bbox.x == 10
|
||
|
|
assert item.bbox.y == 10
|
||
|
|
assert item.cropped_image_path == "crops/1/0.jpg"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_crop_is_uploaded(mock_llm, mock_storage, image_bytes, req):
|
||
|
|
mock_storage.download_bytes = AsyncMock(return_value=image_bytes)
|
||
|
|
mock_llm.chat_vision = AsyncMock(return_value=SAMPLE_QUADS_JSON)
|
||
|
|
mock_storage.upload_bytes = AsyncMock(return_value=None)
|
||
|
|
|
||
|
|
from app.services.image_service import extract_quads
|
||
|
|
await extract_quads(req, mock_llm, mock_storage)
|
||
|
|
|
||
|
|
# upload_bytes called once for the crop
|
||
|
|
mock_storage.upload_bytes.assert_called_once()
|
||
|
|
call_args = mock_storage.upload_bytes.call_args
|
||
|
|
assert call_args.args[1] == "crops/1/0.jpg"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_out_of_bounds_bbox_is_clamped(mock_llm, mock_storage, req):
|
||
|
|
img = _make_test_image_bytes(width=50, height=40)
|
||
|
|
mock_storage.download_bytes = AsyncMock(return_value=img)
|
||
|
|
|
||
|
|
# bbox goes outside image boundary
|
||
|
|
oob_json = json.dumps([{
|
||
|
|
"subject": "test",
|
||
|
|
"predicate": "rel",
|
||
|
|
"object": "obj",
|
||
|
|
"qualifier": None,
|
||
|
|
"bbox": {"x": 30, "y": 20, "w": 100, "h": 100}, # extends beyond 50x40
|
||
|
|
}])
|
||
|
|
mock_llm.chat_vision = AsyncMock(return_value=oob_json)
|
||
|
|
mock_storage.upload_bytes = AsyncMock(return_value=None)
|
||
|
|
|
||
|
|
from app.services.image_service import extract_quads
|
||
|
|
# Should not raise; bbox is clamped
|
||
|
|
result = await extract_quads(req, mock_llm, mock_storage)
|
||
|
|
assert len(result.items) == 1
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_llm_parse_error_raised(mock_llm, mock_storage, image_bytes, req):
|
||
|
|
mock_storage.download_bytes = AsyncMock(return_value=image_bytes)
|
||
|
|
mock_llm.chat_vision = AsyncMock(return_value="bad json {{")
|
||
|
|
|
||
|
|
from app.services.image_service import extract_quads
|
||
|
|
with pytest.raises(LLMParseError):
|
||
|
|
await extract_quads(req, mock_llm, mock_storage)
|