fix(vision): limite di tipo e dimensione sui tre endpoint di upload
/execute, /preview e /reference-images leggevano l'immagine senza limite di tipo o dimensione, e /reference-images scriveva su disco con estensione .png fissa qualunque fosse il tipo reale. Riusa la convenzione già esistente in files.py (ALLOWED_IMAGE_TYPES, validate_file_size) invece di inventarne una seconda: tipo controllato dall'header prima di leggere, dimensione controllata sul corpo letto, un PDF non è qualcosa che un grafo di visione può misurare. save_reference_image ora scrive con l'estensione del tipo reale dell'upload. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014BBnuACZSCJqXrMYC3LUMU
This commit is contained in:
@@ -0,0 +1,174 @@
|
||||
"""I3: the three vision endpoints must not accept an unbounded or wrong-typed
|
||||
upload. `/execute`, `/preview` and `/reference-images` read the whole image
|
||||
into memory with no size cap, and `/reference-images` wrote whatever arrived
|
||||
to disk with a `.png` extension regardless of its real type - a PDF is not
|
||||
something a vision graph can measure.
|
||||
|
||||
Reuses the project's existing convention (`validate_file_size`,
|
||||
`ALLOWED_IMAGE_TYPES` from `files.py`) instead of inventing a second one.
|
||||
"""
|
||||
import io
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
from sqlalchemy import select
|
||||
|
||||
from src.backend.config import settings
|
||||
from src.backend.models.orm.measurement import Measurement
|
||||
from src.backend.models.orm.recipe import RecipeVersion
|
||||
from src.backend.models.orm.task import RecipeSubtask, RecipeTask
|
||||
from src.backend.models.orm.vision import VisionReferenceImage, VisionResult
|
||||
from src.backend.tests.conftest import auth_headers, create_test_recipe
|
||||
|
||||
|
||||
def a_png() -> bytes:
|
||||
buffer = io.BytesIO()
|
||||
Image.new("L", (64, 64), color=0).save(buffer, format="PNG")
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def a_jpeg() -> bytes:
|
||||
buffer = io.BytesIO()
|
||||
Image.new("RGB", (64, 64), color=0).save(buffer, format="JPEG")
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
async def recipe_parts(db_session, recipe_id: int):
|
||||
"""`create_test_recipe` returns only the recipe: here it is broken apart."""
|
||||
version = (await db_session.execute(
|
||||
select(RecipeVersion).where(
|
||||
RecipeVersion.recipe_id == recipe_id,
|
||||
RecipeVersion.is_current == True, # noqa: E712
|
||||
)
|
||||
)).scalar_one()
|
||||
task = (await db_session.execute(
|
||||
select(RecipeTask).where(RecipeTask.version_id == version.id)
|
||||
)).scalars().first()
|
||||
subtask = (await db_session.execute(
|
||||
select(RecipeSubtask).where(RecipeSubtask.task_id == task.id)
|
||||
)).scalars().first()
|
||||
return version, task, subtask
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def worker_says(monkeypatch):
|
||||
def _install(outputs: dict):
|
||||
async def fake_call(image_bytes: bytes, graph: dict) -> dict:
|
||||
return {"outputs": outputs, "failures": [],
|
||||
"engine_version": "c" * 40, "duration_ms": 5}
|
||||
monkeypatch.setattr(
|
||||
"src.backend.services.vision_service._call_worker", fake_call,
|
||||
)
|
||||
return _install
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_rejects_a_non_image_upload(
|
||||
client, db_session, measurement_tec_user, worker_says,
|
||||
):
|
||||
recipe = await create_test_recipe(db_session, measurement_tec_user.id)
|
||||
version, task, subtask = await recipe_parts(db_session, recipe.id)
|
||||
task.task_type = "camera_measure"
|
||||
task.vision_json = {"schema_version": 1, "name": "p", "tools": []}
|
||||
subtask.vision_output = "t1.diameter"
|
||||
await db_session.commit()
|
||||
worker_says({"t1.diameter": float(subtask.nominal)})
|
||||
|
||||
response = await client.post(
|
||||
"/api/vision/execute",
|
||||
headers=auth_headers(measurement_tec_user),
|
||||
files={"image": ("p.pdf", b"%PDF-1.4 not an image", "application/pdf")},
|
||||
data={"task_id": str(task.id)},
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert (await db_session.execute(select(VisionResult))).scalars().all() == []
|
||||
assert (await db_session.execute(select(Measurement))).scalars().all() == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_rejects_an_oversized_upload(
|
||||
client, db_session, measurement_tec_user, worker_says, monkeypatch,
|
||||
):
|
||||
recipe = await create_test_recipe(db_session, measurement_tec_user.id)
|
||||
version, task, subtask = await recipe_parts(db_session, recipe.id)
|
||||
task.task_type = "camera_measure"
|
||||
task.vision_json = {"schema_version": 1, "name": "p", "tools": []}
|
||||
subtask.vision_output = "t1.diameter"
|
||||
await db_session.commit()
|
||||
worker_says({"t1.diameter": float(subtask.nominal)})
|
||||
monkeypatch.setattr(settings, "max_upload_size_mb", 0)
|
||||
|
||||
response = await client.post(
|
||||
"/api/vision/execute",
|
||||
headers=auth_headers(measurement_tec_user),
|
||||
files={"image": ("p.png", a_png(), "image/png")},
|
||||
data={"task_id": str(task.id)},
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert (await db_session.execute(select(VisionResult))).scalars().all() == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preview_rejects_a_non_image_upload(
|
||||
client, db_session, maker_user, worker_says,
|
||||
):
|
||||
recipe = await create_test_recipe(db_session, maker_user.id)
|
||||
version, task, subtask = await recipe_parts(db_session, recipe.id)
|
||||
task.task_type = "camera_measure"
|
||||
task.vision_json = {"schema_version": 1, "name": "p", "tools": []}
|
||||
await db_session.commit()
|
||||
worker_says({"t1.diameter": 10.0})
|
||||
|
||||
response = await client.post(
|
||||
"/api/vision/preview",
|
||||
headers=auth_headers(maker_user),
|
||||
files={"image": ("p.pdf", b"%PDF-1.4 not an image", "application/pdf")},
|
||||
data={"task_id": str(task.id)},
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reference_images_rejects_a_non_image_upload(
|
||||
client, db_session, maker_user,
|
||||
):
|
||||
recipe = await create_test_recipe(db_session, maker_user.id)
|
||||
_, task, _ = await recipe_parts(db_session, recipe.id)
|
||||
await db_session.commit()
|
||||
|
||||
response = await client.post(
|
||||
"/api/vision/reference-images",
|
||||
headers=auth_headers(maker_user),
|
||||
files={"image": ("p.pdf", b"%PDF-1.4 not an image", "application/pdf")},
|
||||
data={"task_id": str(task.id)},
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert (await db_session.execute(
|
||||
select(VisionReferenceImage)
|
||||
)).scalars().all() == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reference_image_keeps_the_extension_matching_its_real_type(
|
||||
client, db_session, maker_user,
|
||||
):
|
||||
"""A jpeg upload must not be written to disk as `<uuid>.png`."""
|
||||
recipe = await create_test_recipe(db_session, maker_user.id)
|
||||
_, task, _ = await recipe_parts(db_session, recipe.id)
|
||||
await db_session.commit()
|
||||
|
||||
response = await client.post(
|
||||
"/api/vision/reference-images",
|
||||
headers=auth_headers(maker_user),
|
||||
files={"image": ("p.jpg", a_jpeg(), "image/jpeg")},
|
||||
data={"task_id": str(task.id)},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
saved = (await db_session.execute(select(VisionReferenceImage))).scalars().one()
|
||||
assert saved.path.endswith(".jpg")
|
||||
assert (settings.upload_path / saved.path).exists()
|
||||
Reference in New Issue
Block a user