feat(US2): image quad extraction — POST /api/v1/image/extract
- app/models/image_models.py: BBox, QuadrupleItem, ImageExtract{Request,Response}
- app/services/image_service.py: download → base64 LLM → bbox clamp → crop upload
- app/routers/image.py: POST /image/extract handler
- tests: 4 service + 3 router tests, 7/7 passing
This commit is contained in:
63
tests/test_image_router.py
Normal file
63
tests/test_image_router.py
Normal file
@@ -0,0 +1,63 @@
|
||||
import json
|
||||
import numpy as np
|
||||
import cv2
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from app.core.exceptions import StorageError
|
||||
|
||||
|
||||
def _make_test_image_bytes() -> bytes:
|
||||
img = np.zeros((80, 100, 3), dtype=np.uint8)
|
||||
_, buf = cv2.imencode(".jpg", img)
|
||||
return buf.tobytes()
|
||||
|
||||
|
||||
SAMPLE_QUADS_JSON = json.dumps([
|
||||
{
|
||||
"subject": "电缆接头",
|
||||
"predicate": "位于",
|
||||
"object": "配电箱左侧",
|
||||
"qualifier": "2024年检修",
|
||||
"bbox": {"x": 5, "y": 5, "w": 20, "h": 15},
|
||||
}
|
||||
])
|
||||
|
||||
|
||||
def test_image_extract_returns_200(client, mock_llm, mock_storage):
|
||||
mock_storage.download_bytes = AsyncMock(return_value=_make_test_image_bytes())
|
||||
mock_llm.chat_vision = AsyncMock(return_value=SAMPLE_QUADS_JSON)
|
||||
mock_storage.upload_bytes = AsyncMock(return_value=None)
|
||||
|
||||
resp = client.post(
|
||||
"/api/v1/image/extract",
|
||||
json={"file_path": "image/test.jpg", "task_id": 1},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert "items" in data
|
||||
assert data["items"][0]["subject"] == "电缆接头"
|
||||
assert data["items"][0]["cropped_image_path"] == "crops/1/0.jpg"
|
||||
|
||||
|
||||
def test_image_extract_llm_parse_error_returns_502(client, mock_llm, mock_storage):
|
||||
mock_storage.download_bytes = AsyncMock(return_value=_make_test_image_bytes())
|
||||
mock_llm.chat_vision = AsyncMock(return_value="not json {{")
|
||||
|
||||
resp = client.post(
|
||||
"/api/v1/image/extract",
|
||||
json={"file_path": "image/test.jpg", "task_id": 1},
|
||||
)
|
||||
assert resp.status_code == 502
|
||||
assert resp.json()["code"] == "LLM_PARSE_ERROR"
|
||||
|
||||
|
||||
def test_image_extract_storage_error_returns_502(client, mock_storage):
|
||||
mock_storage.download_bytes = AsyncMock(side_effect=StorageError("RustFS down"))
|
||||
|
||||
resp = client.post(
|
||||
"/api/v1/image/extract",
|
||||
json={"file_path": "image/test.jpg", "task_id": 1},
|
||||
)
|
||||
assert resp.status_code == 502
|
||||
assert resp.json()["code"] == "STORAGE_ERROR"
|
||||
Reference in New Issue
Block a user