2876c179ac
- 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
64 lines
1.9 KiB
Python
64 lines
1.9 KiB
Python
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"
|