217 lines
8.0 KiB
Python
217 lines
8.0 KiB
Python
"""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
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def redirect_uploads(monkeypatch, tmp_path):
|
|
"""Keep the reference image this file stores out of the real repo tree.
|
|
|
|
The same guard `test_vision_reference.py` carries: without it, the test that
|
|
proves a jpeg is stored with its own extension leaves a real file under
|
|
`uploads/vision/reference/`, so every suite run dirties the working tree.
|
|
|
|
`settings.upload_path` is `parents[2] / self.upload_dir`, and pathlib
|
|
replaces the left operand of `/` when the right one is absolute - so
|
|
redirecting the bare `upload_dir` setting is enough to redirect the
|
|
resolved path too.
|
|
"""
|
|
monkeypatch.setattr(settings, "upload_dir", str(tmp_path))
|
|
|
|
|
|
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, **kwargs) -> 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_a_file_whose_bytes_do_not_match_its_declared_type(
|
|
client, db_session, measurement_tec_user, worker_says,
|
|
):
|
|
"""#15: a text file declared image/png passes the Content-Type check but
|
|
must still not reach the worker - it is refused on its magic bytes."""
|
|
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.png", b"i am not an image", "image/png")},
|
|
data={"task_id": str(task.id)},
|
|
)
|
|
|
|
assert response.status_code == 422
|
|
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()
|