312 lines
11 KiB
Python
312 lines
11 KiB
Python
"""The worker is a thin shell: it receives an image and a graph, and answers
|
|
with numbers."""
|
|
import io
|
|
import json
|
|
|
|
import pytest
|
|
|
|
# Same reasoning as src/vision/tests/test_runner.py: the vision tree (and
|
|
# therefore this worker, which imports src.vision.runner) requires Python
|
|
# >=3.13. This must run before `import numpy` / `from src.vision_worker.main
|
|
# import app` so that `uv run pytest` on the default 3.11 environment reports
|
|
# this module as skipped, with a reason, instead of erroring out of
|
|
# collection.
|
|
pytest.importorskip(
|
|
"visionsuite.task",
|
|
reason=(
|
|
"the vision tree requires Python >=3.13 (vs-task/vs-measure/vs-pm2d "
|
|
"declare it upstream); run it with "
|
|
"`uv run --python 3.13 --extra vision-worker --extra dev pytest "
|
|
"src/vision_worker/tests`"
|
|
),
|
|
)
|
|
|
|
import asyncio
|
|
import time
|
|
|
|
import numpy as np
|
|
from httpx import ASGITransport, AsyncClient
|
|
from PIL import Image
|
|
|
|
from src.vision import runner
|
|
from src.vision_worker import main as worker_main
|
|
from src.vision_worker.main import app
|
|
|
|
|
|
def a_disc_png(radius: int = 60, size: int = 400) -> bytes:
|
|
image = np.zeros((size, size), dtype=np.uint8)
|
|
yy, xx = np.ogrid[:size, :size]
|
|
centre = size // 2
|
|
image[(yy - centre) ** 2 + (xx - centre) ** 2 <= radius ** 2] = 255
|
|
buffer = io.BytesIO()
|
|
Image.fromarray(image).save(buffer, format="PNG")
|
|
return buffer.getvalue()
|
|
|
|
|
|
def a_circle_graph() -> dict:
|
|
from visionsuite.task import Task
|
|
from visionsuite.tools.detect import DetectCircle
|
|
from visionsuite.tools.dimensions import Diameter
|
|
|
|
task = Task(name="prova")
|
|
found = task.add(DetectCircle(name="disco"))
|
|
found.center_row, found.center_col, found.radius = 200.0, 200.0, 60.0
|
|
found.scan_size = 30.0
|
|
quota = task.add(Diameter(name="diametro"))
|
|
quota.nominal_unit = "px" # #46: v2, senza calibrazione la quota resta in px
|
|
quota.link("circle", found, "circle")
|
|
return task.to_dict()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_health_says_which_engine_it_carries():
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://worker") as client:
|
|
response = await client.get("/health")
|
|
|
|
assert response.status_code == 200
|
|
assert len(response.json()["engine_version"]) == 40
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_returns_the_named_outputs():
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://worker") as client:
|
|
response = await client.post(
|
|
"/run",
|
|
files={"image": ("piece.png", a_disc_png(), "image/png")},
|
|
data={"graph": json.dumps(a_circle_graph())},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
body = response.json()
|
|
assert body["failures"] == []
|
|
assert any(key.endswith(".diameter") for key in body["outputs"])
|
|
# Pinned to src/vision/tests/test_runner.py's in-process expectation
|
|
# (same graph, same image, same tolerance) - the whole point of a shared
|
|
# engine_version is that the number does not depend on which path
|
|
# computed it.
|
|
diameter = next(v for k, v in body["outputs"].items() if k.endswith(".diameter"))
|
|
assert diameter == pytest.approx(120.0, abs=4.0)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_malformed_graph_is_refused_not_crashed():
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://worker") as client:
|
|
response = await client.post(
|
|
"/run",
|
|
files={"image": ("piece.png", a_disc_png(), "image/png")},
|
|
data={"graph": json.dumps({"schema_version": 99, "tools": []})},
|
|
)
|
|
|
|
assert response.status_code == 422
|
|
assert "99" in response.json()["detail"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# #15: a non-image upload, a graph that is not a dict, or a tool with no
|
|
# "kind" used to reach the worker as an unhandled exception - a naked 500,
|
|
# and then a 502 with no reason once the backend forwarded it.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_non_image_upload_is_refused_not_crashed():
|
|
"""A text file declared image/png: Image.open() must not escape as 500."""
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://worker") as client:
|
|
response = await client.post(
|
|
"/run",
|
|
files={"image": ("piece.png", b"i am not an image", "image/png")},
|
|
data={"graph": json.dumps(a_circle_graph())},
|
|
)
|
|
|
|
assert response.status_code == 422
|
|
assert "not decodable" in response.json()["detail"].lower()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_graph_that_is_not_an_object_is_refused():
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://worker") as client:
|
|
response = await client.post(
|
|
"/run",
|
|
files={"image": ("piece.png", a_disc_png(), "image/png")},
|
|
data={"graph": "[]"},
|
|
)
|
|
|
|
assert response.status_code == 422
|
|
assert "object" in response.json()["detail"].lower()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_tool_without_a_kind_is_refused_not_crashed():
|
|
graph = {"schema_version": 2, "name": "p", "tools": [{"id": "a", "name": "x"}]}
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://worker") as client:
|
|
response = await client.post(
|
|
"/run",
|
|
files={"image": ("piece.png", a_disc_png(), "image/png")},
|
|
data={"graph": json.dumps(graph)},
|
|
)
|
|
|
|
assert response.status_code == 422
|
|
assert "malformed graph" in response.json()["detail"].lower()
|
|
assert "tipo sconosciuto" in response.json()["detail"].lower() # #46: messaggio v2
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# #34: /run used to execute PIL decode + run_graph synchronously inside the
|
|
# event loop, with no timeout - a slow graph blocked every other request,
|
|
# /health included, and there was no cap on how many ran at once.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_run_past_its_budget_answers_504_and_health_still_answers(
|
|
monkeypatch,
|
|
):
|
|
monkeypatch.setattr(worker_main, "_RUN_TIMEOUT_S", 0.1)
|
|
|
|
def slow_run_graph(frame, graph, **kwargs):
|
|
time.sleep(0.5) # off the event loop - it runs inside asyncio.to_thread
|
|
return runner.RunOutcome(engine_version="a" * 40)
|
|
|
|
monkeypatch.setattr(worker_main, "run_graph", slow_run_graph)
|
|
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://worker") as client:
|
|
run_task = asyncio.create_task(client.post(
|
|
"/run",
|
|
files={"image": ("piece.png", a_disc_png(), "image/png")},
|
|
data={"graph": json.dumps(a_circle_graph())},
|
|
))
|
|
await asyncio.sleep(0.05) # let /run start and enter the thread
|
|
|
|
health_response = await client.get("/health")
|
|
assert health_response.status_code == 200
|
|
|
|
run_response = await run_task
|
|
|
|
assert run_response.status_code == 504
|
|
assert "0.1" in run_response.json()["detail"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_third_concurrent_run_is_refused_not_queued(monkeypatch):
|
|
monkeypatch.setattr(worker_main, "_MAX_CONCURRENCY", 1)
|
|
monkeypatch.setattr(worker_main, "_run_slots", asyncio.Semaphore(1))
|
|
monkeypatch.setattr(worker_main, "_RUN_TIMEOUT_S", 2.0)
|
|
|
|
def slow_run_graph(frame, graph, **kwargs):
|
|
time.sleep(0.3)
|
|
return runner.RunOutcome(engine_version="a" * 40)
|
|
|
|
monkeypatch.setattr(worker_main, "run_graph", slow_run_graph)
|
|
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://worker") as client:
|
|
first = asyncio.create_task(client.post(
|
|
"/run",
|
|
files={"image": ("piece.png", a_disc_png(), "image/png")},
|
|
data={"graph": json.dumps(a_circle_graph())},
|
|
))
|
|
await asyncio.sleep(0.05) # let the first request take the only slot
|
|
|
|
second_response = await client.post(
|
|
"/run",
|
|
files={"image": ("piece.png", a_disc_png(), "image/png")},
|
|
data={"graph": json.dumps(a_circle_graph())},
|
|
)
|
|
|
|
first_response = await first
|
|
|
|
assert second_response.status_code == 503
|
|
assert first_response.status_code == 200
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# #46: recipe_dir, calibration, overlay, /validate.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def _post(path: str, **kwargs):
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://worker") as client:
|
|
return await client.post(path, **kwargs)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_returns_overlay_and_outcome():
|
|
response = await _post(
|
|
"/run",
|
|
files={"image": ("piece.png", a_disc_png(), "image/png")},
|
|
data={"graph": json.dumps(a_circle_graph())},
|
|
)
|
|
|
|
body = response.json()
|
|
assert response.status_code == 200
|
|
assert body["overlay"] and "tool_id" in body["overlay"][0]
|
|
assert body["outcome"] in ("pass", "warning", "fail")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_accepts_a_calibration():
|
|
graph = a_circle_graph()
|
|
for tool in graph["tools"]:
|
|
if tool["kind"] == "Diameter":
|
|
tool["values"]["nominal_unit"] = "mm"
|
|
response = await _post(
|
|
"/run",
|
|
files={"image": ("piece.png", a_disc_png(), "image/png")},
|
|
data={"graph": json.dumps(graph),
|
|
"calibration": json.dumps({"px_per_mm": 2.0})},
|
|
)
|
|
|
|
body = response.json()
|
|
assert response.status_code == 200, body
|
|
assert body["failures"] == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("bad", ["../etc", "a/../../b", "/abs/path", "C:/x/y"])
|
|
async def test_a_recipe_dir_outside_the_root_is_refused(bad, monkeypatch, tmp_path):
|
|
monkeypatch.setenv("VISION_RECIPES_ROOT", str(tmp_path))
|
|
response = await _post(
|
|
"/run",
|
|
files={"image": ("piece.png", a_disc_png(), "image/png")},
|
|
data={"graph": json.dumps(a_circle_graph()), "recipe_dir": bad},
|
|
)
|
|
assert response.status_code == 422
|
|
|
|
response = await _post(
|
|
"/validate", data={"graph": json.dumps(a_circle_graph()), "recipe_dir": bad}
|
|
)
|
|
assert response.status_code == 422
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_validate_ok_and_with_a_missing_file(monkeypatch, tmp_path):
|
|
monkeypatch.setenv("VISION_RECIPES_ROOT", str(tmp_path))
|
|
response = await _post("/validate", data={"graph": json.dumps(a_circle_graph())})
|
|
body = response.json()
|
|
assert response.status_code == 200 and body["ok"] is True
|
|
assert any(o["key"].endswith(".diameter") for o in body["outputs"])
|
|
|
|
(tmp_path / "r1").mkdir()
|
|
graph = {
|
|
"schema_version": 2, "name": "d",
|
|
"meta": {"dxf": {"d": {"path": "mancante.dxf"}}},
|
|
"tools": [{"id": "t", "name": "c", "kind": "CompareDxf",
|
|
"values": {"dxf": "d", "mm_per_px": 0.25,
|
|
"band_plus": 8.0, "band_minus": 8.0},
|
|
"links": {}}],
|
|
}
|
|
response = await _post(
|
|
"/validate", data={"graph": json.dumps(graph), "recipe_dir": "r1"}
|
|
)
|
|
body = response.json()
|
|
assert response.status_code == 200
|
|
assert body["ok"] is False and body["errors"] and body["outputs"] == []
|