387 lines
16 KiB
Python
387 lines
16 KiB
Python
"""
|
|
Contract tests for the recognition server API.
|
|
|
|
A fake engine stands in for the model, so these run without weights. They defend
|
|
the wire contracts a client depends on: response shapes, the shared error shape,
|
|
auth, limits, and the WebSocket protocol (hello, frames, control messages,
|
|
recoverable errors).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import json
|
|
import os
|
|
import sys
|
|
import unittest
|
|
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
if current_dir not in sys.path:
|
|
sys.path.insert(0, current_dir)
|
|
|
|
import cv2 # noqa: E402
|
|
import numpy as np # noqa: E402
|
|
|
|
from app import ServerSettings, create_app # noqa: E402
|
|
from engine import EngineSettings, RecognitionEngine # noqa: E402
|
|
from test_live import FakeDetector, FakeInventory, FakeMatcher # noqa: E402
|
|
|
|
from fastapi.testclient import TestClient # noqa: E402
|
|
|
|
|
|
def make_scans(engine):
|
|
"""Legacy session model wired to the fake engine (the real factory is exercised
|
|
against the real model in test_real_model)."""
|
|
from observe import PerceptionAdapter, PerceptionService, ScanService
|
|
|
|
def detect(frame_bytes):
|
|
image = engine.decode_image(frame_bytes)
|
|
boxes = engine.detect(image)
|
|
return {"width": image.shape[1], "height": image.shape[0], "boxes": boxes,
|
|
"empty_slots": [], "quality": {"decoded": True}}
|
|
|
|
def classify(box):
|
|
scores = engine.score_crop(engine.crop(engine._last_image, box["box"]))
|
|
return scores
|
|
|
|
# keep the frame for the classify callback
|
|
def detect_with_frame(frame_bytes):
|
|
image = engine.decode_image(frame_bytes)
|
|
engine._last_image = image
|
|
boxes = engine.detect(image)
|
|
return {"width": image.shape[1], "height": image.shape[0], "boxes": boxes,
|
|
"empty_slots": [], "quality": {"decoded": True}}
|
|
|
|
adapter = PerceptionAdapter(detect_fn=detect_with_frame, classify_fn=classify,
|
|
model_version="fake", matcher_method="hybrid")
|
|
return ScanService(PerceptionService(adapter))
|
|
|
|
|
|
def make_app(*, token=None, max_sessions=4):
|
|
engine = RecognitionEngine(
|
|
FakeDetector(), FakeMatcher(), FakeInventory(),
|
|
settings=EngineSettings(classify_workers=2),
|
|
)
|
|
settings = ServerSettings(token=token, max_sessions=max_sessions)
|
|
app = create_app(engine, settings, web_dir="/nonexistent", scans=make_scans(engine))
|
|
return TestClient(app), engine
|
|
|
|
|
|
def jpeg_bytes(width=640, height=480):
|
|
ok, buffer = cv2.imencode(".jpg", np.zeros((height, width, 3), dtype="uint8"))
|
|
return buffer.tobytes()
|
|
|
|
|
|
def frame_payload(frame_id, payload=None):
|
|
payload = payload if payload is not None else jpeg_bytes()
|
|
return frame_id.to_bytes(4, "big") + payload
|
|
|
|
|
|
class RestContractTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.client, self.engine = make_app()
|
|
|
|
def tearDown(self):
|
|
self.engine.shutdown()
|
|
|
|
def test_health_needs_no_token_and_reports_model(self):
|
|
response = self.client.get("/v1/health")
|
|
self.assertEqual(response.status_code, 200)
|
|
body = response.json()
|
|
self.assertEqual(body["status"], "ok")
|
|
self.assertIn("version", body["model"])
|
|
self.assertEqual(body["catalog"]["sku_count"], 2)
|
|
|
|
def test_catalog_shape_has_no_local_paths(self):
|
|
body = self.client.get("/v1/catalog").json()
|
|
product = body["products"][0]
|
|
self.assertEqual(set(product), {
|
|
"sku_id", "product_name", "price", "number_in", "number_out", "number_on",
|
|
"receipt_date", "expiration_date", "image_url",
|
|
})
|
|
|
|
def test_catalog_image_404_for_unknown_and_traversal(self):
|
|
response = self.client.get("/v1/catalog/nope/image")
|
|
self.assertEqual(response.status_code, 404)
|
|
self.assertEqual(response.json()["error"]["code"], "NOT_FOUND")
|
|
# A traversal-looking id is rejected before it can reach the filesystem
|
|
# (either by routing or by the inventory lookup).
|
|
for sku in ("..%2F..%2Fstock.csv", "../../etc/passwd", "sku_a%2F..%2F..%2Fstock.csv"):
|
|
response = self.client.get(f"/v1/catalog/{sku}/image")
|
|
self.assertIn(response.status_code, (400, 404), sku)
|
|
|
|
def test_stock_update_validates_and_returns_product(self):
|
|
before = self.client.get("/v1/catalog").json()["products"]
|
|
before_on = next(p["number_on"] for p in before if p["sku_id"] == "sku_a")
|
|
response = self.client.post("/v1/catalog/sku_a/stock", json={"delta": -1})
|
|
self.assertEqual(response.status_code, 200)
|
|
self.assertEqual(response.json()["product"]["number_on"], before_on - 1)
|
|
self.assertEqual(self.client.post("/v1/catalog/sku_a/stock", json={"delta": 0}).status_code, 422)
|
|
self.assertEqual(self.client.post("/v1/catalog/sku_a/stock", json={"delta": 500}).status_code, 422)
|
|
self.assertEqual(self.client.post("/v1/catalog/ghost/stock", json={"delta": -1}).status_code, 404)
|
|
|
|
def test_recognize_accepts_raw_bytes(self):
|
|
response = self.client.post(
|
|
"/v1/recognize?classify=true", content=jpeg_bytes(),
|
|
headers={"Content-Type": "image/jpeg"},
|
|
)
|
|
self.assertEqual(response.status_code, 200)
|
|
body = response.json()
|
|
self.assertTrue(body["image_id"].startswith("img_"))
|
|
self.assertEqual(body["coordinate_convention"], "pixel-xyxy-origin-top-left")
|
|
self.assertEqual(len(body["detections"]), 1)
|
|
self.assertIn(body["detections"][0]["identity"]["state"], ("recognized", "unconfirmed"))
|
|
for key in ("width", "height", "catalog_version", "model_version", "summary", "timing_ms"):
|
|
self.assertIn(key, body)
|
|
|
|
def test_recognize_accepts_multipart(self):
|
|
response = self.client.post(
|
|
"/v1/recognize?classify=false",
|
|
files={"file": ("shelf.jpg", jpeg_bytes(), "image/jpeg")},
|
|
)
|
|
self.assertEqual(response.status_code, 200)
|
|
self.assertEqual(response.json()["detections"][0]["identity"]["state"], "detected")
|
|
|
|
def test_recognize_rejects_bad_input(self):
|
|
self.assertEqual(self.client.post("/v1/recognize", content=b"").status_code, 400)
|
|
response = self.client.post(
|
|
"/v1/recognize", content=b"not an image", headers={"Content-Type": "image/jpeg"}
|
|
)
|
|
self.assertEqual(response.status_code, 422)
|
|
self.assertEqual(response.json()["error"]["code"], "UNDECODABLE_IMAGE")
|
|
response = self.client.post(
|
|
"/v1/recognize", content=b"{}", headers={"Content-Type": "application/json"}
|
|
)
|
|
self.assertEqual(response.status_code, 415)
|
|
self.assertEqual(response.json()["error"]["code"], "UNSUPPORTED_MEDIA_TYPE")
|
|
|
|
def test_a_lying_content_length_is_rejected_while_streaming(self):
|
|
"""The ASGI receive stops once the body passes the cap, before decode."""
|
|
import anyio
|
|
|
|
from transport import BodyLimit
|
|
|
|
cap = 32
|
|
seen = {"bytes": 0}
|
|
|
|
async def app(scope, receive, send):
|
|
while True:
|
|
message = await receive()
|
|
if message["type"] != "http.request":
|
|
continue
|
|
seen["bytes"] += len(message.get("body") or b"")
|
|
if not message.get("more_body"):
|
|
break
|
|
await send({"type": "http.response.start", "status": 200, "headers": []})
|
|
await send({"type": "http.response.body", "body": b"ok"})
|
|
|
|
chunks = [b"x" * 20, b"y" * 20, b"z" * 20]
|
|
|
|
async def receive():
|
|
body = chunks.pop(0) if chunks else b""
|
|
return {"type": "http.request", "body": body, "more_body": bool(chunks)}
|
|
|
|
sent = []
|
|
|
|
async def send(message):
|
|
sent.append(message)
|
|
|
|
async def run():
|
|
await BodyLimit(app, cap)(
|
|
{"type": "http"}, receive, send
|
|
)
|
|
|
|
anyio.run(run)
|
|
self.assertLessEqual(seen["bytes"], cap + 20)
|
|
self.assertEqual(sent[0]["status"], 413)
|
|
self.assertIn(b"PAYLOAD_TOO_LARGE", sent[1]["body"])
|
|
|
|
def test_classify_and_inspect_roundtrip(self):
|
|
recognized = self.client.post(
|
|
"/v1/recognize?classify=false", content=jpeg_bytes(),
|
|
headers={"Content-Type": "image/jpeg"},
|
|
).json()
|
|
image_id = recognized["image_id"]
|
|
box = recognized["detections"][0]["box"]
|
|
|
|
classified = self.client.post(
|
|
f"/v1/images/{image_id}/classify", json={"method": "hybrid"}
|
|
)
|
|
self.assertEqual(classified.status_code, 200)
|
|
self.assertEqual(len(classified.json()["identities"]), 1)
|
|
|
|
inspected = self.client.post(
|
|
f"/v1/images/{image_id}/inspect", json={"box": box, "method": "hybrid"}
|
|
)
|
|
self.assertEqual(inspected.status_code, 200)
|
|
body = inspected.json()
|
|
self.assertIn("identity", body)
|
|
self.assertTrue(body["crop_jpeg_b64"].startswith("data:image/jpeg;base64,"))
|
|
|
|
def test_unknown_method_is_a_400(self):
|
|
response = self.client.post(
|
|
"/v1/recognize?method=nope", content=jpeg_bytes(),
|
|
headers={"Content-Type": "image/jpeg"},
|
|
)
|
|
self.assertEqual(response.status_code, 400)
|
|
self.assertEqual(response.json()["error"]["code"], "BAD_REQUEST")
|
|
|
|
def test_expired_image_id(self):
|
|
response = self.client.post("/v1/images/img_missing/classify", json={"method": "hybrid"})
|
|
self.assertEqual(response.status_code, 404)
|
|
self.assertEqual(response.json()["error"]["code"], "IMAGE_EXPIRED")
|
|
|
|
def test_validation_errors_use_the_shared_shape(self):
|
|
response = self.client.post("/v1/images/img_x/inspect", json={"box": [1, 2, 3]})
|
|
self.assertEqual(response.status_code, 422)
|
|
self.assertIn("error", response.json())
|
|
|
|
|
|
class AuthTests(unittest.TestCase):
|
|
def test_token_is_required_when_configured(self):
|
|
client, engine = make_app(token="s3cret")
|
|
try:
|
|
self.assertEqual(client.get("/v1/health").status_code, 200)
|
|
self.assertEqual(client.get("/v1/catalog").status_code, 401)
|
|
self.assertEqual(client.get("/v1/catalog", headers={"Authorization": "Bearer wrong"}).status_code, 401)
|
|
self.assertEqual(client.get("/v1/catalog", headers={"Authorization": "Bearer s3cret"}).status_code, 200)
|
|
finally:
|
|
engine.shutdown()
|
|
|
|
|
|
class WebSocketProtocolTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.client, self.engine = make_app()
|
|
|
|
def tearDown(self):
|
|
self.engine.shutdown()
|
|
|
|
def test_hello_frame_result_and_control_messages(self):
|
|
with self.client.websocket_connect("/v1/stream") as socket:
|
|
hello = socket.receive_json()
|
|
self.assertEqual(hello["type"], "hello")
|
|
self.assertEqual(hello["protocol"], 1)
|
|
self.assertIn("session_id", hello)
|
|
self.assertIn("config", hello)
|
|
self.assertEqual(hello["limits"]["max_frame_bytes"], 4 * 1024 * 1024)
|
|
|
|
socket.send_json({"type": "ping", "t": 7})
|
|
self.assertEqual(socket.receive_json(), {"type": "pong", "t": 7})
|
|
|
|
socket.send_json({"type": "config", "method": "orb"})
|
|
ack = socket.receive_json()
|
|
self.assertEqual(ack["type"], "config_ack")
|
|
self.assertEqual(ack["config"]["method"], "orb")
|
|
|
|
socket.send_bytes(frame_payload(11))
|
|
result = socket.receive_json()
|
|
self.assertEqual(result["type"], "result")
|
|
self.assertEqual(result["frame_id"], 11)
|
|
self.assertEqual(result["seq"], 1)
|
|
self.assertIn("tracks", result)
|
|
self.assertIn("summary", result)
|
|
self.assertIn("quality", result)
|
|
|
|
socket.send_json({"type": "reset"})
|
|
self.assertEqual(socket.receive_json()["type"], "reset_ack")
|
|
|
|
socket.send_json({"type": "not-a-thing"})
|
|
error = socket.receive_json()
|
|
self.assertEqual(error["code"], "BAD_MESSAGE")
|
|
|
|
def test_short_frame_is_recoverable(self):
|
|
with self.client.websocket_connect("/v1/stream") as socket:
|
|
socket.receive_json() # hello
|
|
socket.send_bytes(b"\x00\x01")
|
|
error = socket.receive_json()
|
|
self.assertEqual(error["code"], "BAD_FRAME")
|
|
# the connection is still usable
|
|
socket.send_bytes(frame_payload(1))
|
|
self.assertEqual(socket.receive_json()["type"], "result")
|
|
|
|
def test_undecodable_frame_is_reported_and_keeps_the_socket(self):
|
|
with self.client.websocket_connect("/v1/stream") as socket:
|
|
socket.receive_json()
|
|
socket.send_bytes(frame_payload(5, b"not an image"))
|
|
error = socket.receive_json()
|
|
self.assertEqual(error["code"], "UNDECODABLE_IMAGE")
|
|
self.assertEqual(error["frame_id"], 5)
|
|
socket.send_bytes(frame_payload(6))
|
|
self.assertEqual(socket.receive_json()["frame_id"], 6)
|
|
|
|
def test_inspect_over_the_socket(self):
|
|
with self.client.websocket_connect("/v1/stream") as socket:
|
|
socket.receive_json()
|
|
socket.send_bytes(frame_payload(1))
|
|
result = socket.receive_json()
|
|
track_id = result["tracks"][0]["track_id"]
|
|
socket.send_json({"type": "inspect", "track_id": track_id, "request_id": "r1"})
|
|
reply = socket.receive_json()
|
|
self.assertEqual(reply["type"], "inspect_result")
|
|
self.assertEqual(reply["request_id"], "r1")
|
|
self.assertIn("identity", reply)
|
|
|
|
socket.send_json({"type": "inspect", "track_id": 999999, "request_id": "r2"})
|
|
missing = socket.receive_json()
|
|
self.assertEqual(missing["code"], "TRACK_NOT_IN_VIEW")
|
|
|
|
def test_token_is_enforced_on_the_socket(self):
|
|
"""Header or the first post-hello auth frame. A query token is rejected."""
|
|
client, engine = make_app(token="s3cret")
|
|
try:
|
|
with self.assertRaises(Exception):
|
|
with client.websocket_connect("/v1/stream") as socket:
|
|
self.assertEqual(socket.receive_json()["type"], "hello")
|
|
socket.send_json({"type": "ping", "t": 1})
|
|
socket.receive_json()
|
|
with self.assertRaises(Exception):
|
|
with client.websocket_connect("/v1/stream?token=s3cret") as socket:
|
|
socket.receive_json()
|
|
with client.websocket_connect(
|
|
"/v1/stream", headers={"Authorization": "Bearer s3cret"}
|
|
) as socket:
|
|
self.assertEqual(socket.receive_json()["type"], "hello")
|
|
with client.websocket_connect("/v1/stream") as socket:
|
|
self.assertEqual(socket.receive_json()["type"], "hello")
|
|
socket.send_json({"type": "auth", "token": "s3cret"})
|
|
socket.send_json({"type": "ping", "t": 1})
|
|
self.assertEqual(socket.receive_json(), {"type": "pong", "t": 1})
|
|
finally:
|
|
engine.shutdown()
|
|
|
|
def test_session_limit_closes_extra_clients(self):
|
|
client, engine = make_app(max_sessions=1)
|
|
try:
|
|
with client.websocket_connect("/v1/stream") as first:
|
|
first.receive_json()
|
|
with self.assertRaises(Exception):
|
|
with client.websocket_connect("/v1/stream") as second:
|
|
second.receive_json()
|
|
finally:
|
|
engine.shutdown()
|
|
|
|
|
|
class LegacyRouteTests(unittest.TestCase):
|
|
def test_legacy_sessions_still_respond(self):
|
|
"""The earlier session model stays mounted; it is backed by the engine now."""
|
|
client, engine = make_app()
|
|
try:
|
|
started = client.post("/api/sessions", json={"zone_id": "z1", "catalog_version": "v1"})
|
|
self.assertEqual(started.status_code, 200)
|
|
session_id = started.json()["session"]["session_id"]
|
|
frames = client.post(
|
|
f"/api/sessions/{session_id}/frames",
|
|
json={
|
|
"frame_id": "f1",
|
|
"sequence": 1,
|
|
"frame_bytes_b64": base64.b64encode(jpeg_bytes()).decode(),
|
|
},
|
|
)
|
|
self.assertEqual(frames.status_code, 200)
|
|
self.assertIn("observation", frames.json())
|
|
finally:
|
|
engine.shutdown()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|