tngo
aaed5d388f
Replaces the DINO gallery as the recognition path with the measured design from docs/research/2026-09-28-recognition-architecture/DECISION.md: RootSIFT over a per-SKU reference pyramid, similarity-constrained RANSAC and a hue check, fused with an OCR text lane that vetoes a match whose package text names a different catalog row and answers on its own when geometry finds nothing. - identity_lanes.py: LocalLane, TextLane, DecisionPolicy. The text lane builds each row's signature from the catalog name plus its packshot OCR, so a row with no enrolled image is still answerable. - engine.py: every backend groups scores by canonical SKU, so no view-suffixed id reaches the wire and a name and price always resolve. decide_crop fuses the lanes; score_crop_lanes feeds the live path. Both lanes are built in warm_up (non-fatally) so the first frame is not the slow one, and the text lane shares the engine's OCR lock. - live.py: decides per crop on the lane's own rank order. The lane's order comes from its rank key and its value is only the gate, so re-sorting by value named a different SKU than the lane chose; averaging in an empty observation also diluted a good confidence below the accept line. - identity.py: a `local` score scale whose gate is one absolute line, because the value is already a margin over the runner-up. - app.py, run-demo-server.sh: `--identity local` is the default; the demo server no longer starts on the leaked `stock-lopia` catalog. - eval_live.py: the replay vocabulary accepts `local`, so a recording from the shipped server can be scored. - README.md: documents the three backends and what each is for. Measured on the corrected image6 labels at live resolution: 36/61 in-catalog facings named at 95% precision, 0 of 48 unknowns accepted, 11 labels sourced from the text lane alone. Through the real server: /v1/recognize 90% precision, /v1/stream 100% precision, no view-suffixed ids on either.
529 lines
24 KiB
Python
529 lines
24 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 SmallBoxDetector # noqa: E402
|
|
from test_live import FakeDetector, FakeInventory, FakeMatcher # noqa: E402
|
|
|
|
from fastapi.testclient import TestClient # noqa: E402
|
|
|
|
|
|
|
|
def make_app(*, token=None, max_sessions=4, detector=None):
|
|
# The fake matcher is what these tests exercise, so the backend is pinned to
|
|
# `matcher`; the shipped default is the `local` lane, which reads real
|
|
# reference images and would ignore the fake.
|
|
engine = RecognitionEngine(
|
|
detector or FakeDetector(), FakeMatcher(), FakeInventory(),
|
|
settings=EngineSettings(classify_workers=2, identity_backend="matcher"),
|
|
)
|
|
settings = ServerSettings(token=token, max_sessions=max_sessions)
|
|
app = create_app(engine, settings, web_dir="/nonexistent")
|
|
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_overflow_after_response_start_sends_no_second_start(self):
|
|
"""A body that overflows after headers began must not start again."""
|
|
import anyio
|
|
|
|
from transport import BodyLimit
|
|
|
|
async def app(scope, receive, send):
|
|
await send({"type": "http.response.start", "status": 200, "headers": []})
|
|
await receive()
|
|
await send({"type": "http.response.body", "body": b"ok"})
|
|
|
|
sent = []
|
|
|
|
async def receive():
|
|
return {"type": "http.request", "body": b"x" * 40, "more_body": False}
|
|
|
|
async def send(message):
|
|
sent.append(message)
|
|
|
|
anyio.run(BodyLimit(app, 8), {"type": "http"}, receive, send)
|
|
starts = [message for message in sent if message.get("type") == "http.response.start"]
|
|
self.assertEqual(len(starts), 1)
|
|
self.assertEqual(starts[0]["status"], 200)
|
|
self.assertFalse(any(message.get("status") == 413 for message in sent))
|
|
|
|
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 crop_app(self):
|
|
"""A shelf whose only facing is too small for the frame crop to decide."""
|
|
client, engine = make_app(detector=SmallBoxDetector())
|
|
self.addCleanup(engine.shutdown)
|
|
return client
|
|
|
|
def test_capabilities_turn_on_the_crop_lane(self):
|
|
with self.client.websocket_connect("/v1/stream") as socket:
|
|
hello = socket.receive_json()
|
|
self.assertIn("max_crop_bytes", hello["limits"])
|
|
socket.send_json({"type": "capabilities", "crop": True, "max_crop_bytes": 65536})
|
|
ack = socket.receive_json()
|
|
self.assertEqual(ack["type"], "capabilities_ack")
|
|
self.assertEqual(ack["capabilities"], {"crop": True, "max_crop_bytes": 65536})
|
|
|
|
def test_a_capability_the_server_cannot_honour_is_refused(self):
|
|
"""A half-stated or malformed capability turns nothing on."""
|
|
for message in (
|
|
{"type": "capabilities", "max_crop_bytes": 65536},
|
|
{"type": "capabilities", "crop": False, "max_crop_bytes": 65536},
|
|
{"type": "capabilities", "crop": "true", "max_crop_bytes": 65536},
|
|
{"type": "capabilities", "crop": True},
|
|
{"type": "capabilities", "crop": True, "max_crop_bytes": 0},
|
|
{"type": "capabilities", "crop": True, "max_crop_bytes": -4096},
|
|
{"type": "capabilities", "crop": True, "max_crop_bytes": True},
|
|
{"type": "capabilities", "crop": True, "max_crop_bytes": "65536"},
|
|
{"type": "capabilities", "crop": True, "max_crop_bytes": 64 * 1024 * 1024},
|
|
):
|
|
with self.subTest(capability=message):
|
|
with self.crop_app().websocket_connect("/v1/stream") as socket:
|
|
socket.receive_json() # hello
|
|
socket.send_json(message)
|
|
self.assertEqual(socket.receive_json()["code"], "BAD_MESSAGE")
|
|
# A refused capability leaves this session on protocol v1: the
|
|
# facing it cannot report is never asked about.
|
|
for frame_id in (1, 2, 3):
|
|
socket.send_bytes(frame_payload(frame_id))
|
|
result = socket.receive_json()
|
|
self.assertEqual(result["stats"]["crops"]["requested"], 0)
|
|
for track in result["tracks"]:
|
|
self.assertNotIn("crop_request", track)
|
|
|
|
def test_a_v1_client_is_never_asked_for_a_crop(self):
|
|
with self.crop_app().websocket_connect("/v1/stream") as socket:
|
|
socket.receive_json() # hello
|
|
for frame_id in (1, 2, 3):
|
|
socket.send_bytes(frame_payload(frame_id))
|
|
result = socket.receive_json()
|
|
for track in result["tracks"]:
|
|
self.assertNotIn("crop_request", track)
|
|
self.assertEqual(result["stats"]["crops"]["requested"], 0)
|
|
|
|
def test_a_crop_round_trip_labels_the_track(self):
|
|
with self.crop_app().websocket_connect("/v1/stream") as socket:
|
|
socket.receive_json() # hello
|
|
socket.send_json({"type": "capabilities", "crop": True, "max_crop_bytes": 131072})
|
|
self.assertEqual(socket.receive_json()["type"], "capabilities_ack")
|
|
socket.send_bytes(frame_payload(1))
|
|
socket.receive_json()
|
|
socket.send_bytes(frame_payload(2))
|
|
result = socket.receive_json()
|
|
request = result["tracks"][0].get("crop_request")
|
|
self.assertIsNotNone(request, "a facing too small for the frame crop must be asked about")
|
|
self.assertEqual(request["frame_id"], 2)
|
|
|
|
socket.send_json({
|
|
"type": "crop",
|
|
"request_id": request["request_id"],
|
|
"track_id": request["track_id"],
|
|
"frame_id": request["frame_id"],
|
|
"jpeg": base64.b64encode(jpeg_bytes(80, 120)).decode("ascii"),
|
|
})
|
|
self.assertEqual(socket.receive_json()["type"], "crop_ack")
|
|
socket.send_bytes(frame_payload(3))
|
|
settled = socket.receive_json()
|
|
self.assertEqual(settled["stats"]["crops"]["accepted"], 1)
|
|
identities = [t["identity"] for t in settled["tracks"]]
|
|
self.assertTrue(any(i["source"] == "crop" for i in identities))
|
|
|
|
def test_a_crop_that_does_not_match_its_request_is_refused(self):
|
|
with self.crop_app().websocket_connect("/v1/stream") as socket:
|
|
socket.receive_json() # hello
|
|
socket.send_json({"type": "capabilities", "crop": True, "max_crop_bytes": 131072})
|
|
socket.receive_json()
|
|
socket.send_bytes(frame_payload(1))
|
|
socket.receive_json()
|
|
socket.send_bytes(frame_payload(2))
|
|
request = socket.receive_json()["tracks"][0]["crop_request"]
|
|
|
|
socket.send_json({
|
|
"type": "crop",
|
|
"request_id": request["request_id"],
|
|
"track_id": request["track_id"],
|
|
"frame_id": 4242,
|
|
"jpeg": base64.b64encode(jpeg_bytes(80, 120)).decode("ascii"),
|
|
})
|
|
refused = socket.receive_json()
|
|
self.assertEqual(refused["type"], "crop_error")
|
|
self.assertEqual(refused["code"], "CROP_MISMATCH")
|
|
socket.send_bytes(frame_payload(3))
|
|
self.assertEqual(socket.receive_json()["stats"]["crops"]["accepted"], 0)
|
|
|
|
def test_a_crop_that_is_not_an_image_is_refused(self):
|
|
with self.crop_app().websocket_connect("/v1/stream") as socket:
|
|
socket.receive_json() # hello
|
|
socket.send_json({"type": "capabilities", "crop": True, "max_crop_bytes": 131072})
|
|
socket.receive_json()
|
|
socket.send_bytes(frame_payload(1))
|
|
socket.receive_json()
|
|
socket.send_bytes(frame_payload(2))
|
|
request = socket.receive_json()["tracks"][0]["crop_request"]
|
|
|
|
socket.send_json({
|
|
"type": "crop",
|
|
"request_id": request["request_id"],
|
|
"track_id": request["track_id"],
|
|
"frame_id": request["frame_id"],
|
|
"jpeg": base64.b64encode(b"not an image").decode("ascii"),
|
|
})
|
|
self.assertEqual(socket.receive_json()["code"], "CROP_UNDECODABLE")
|
|
|
|
def test_a_small_jpeg_claiming_huge_pixels_is_refused_before_decode(self):
|
|
with self.crop_app().websocket_connect("/v1/stream") as socket:
|
|
socket.receive_json()
|
|
socket.send_json({"type": "capabilities", "crop": True, "max_crop_bytes": 131072})
|
|
socket.receive_json()
|
|
socket.send_bytes(frame_payload(1))
|
|
socket.receive_json()
|
|
socket.send_bytes(frame_payload(2))
|
|
request = socket.receive_json()["tracks"][0]["crop_request"]
|
|
jpeg = bytearray(jpeg_bytes(80, 120))
|
|
sof = jpeg.index(b"\xff\xc0")
|
|
jpeg[sof + 5:sof + 9] = (16384).to_bytes(2, "big") * 2
|
|
socket.send_json({
|
|
"type": "crop", "request_id": request["request_id"],
|
|
"track_id": request["track_id"], "frame_id": request["frame_id"],
|
|
"jpeg": base64.b64encode(jpeg).decode("ascii"),
|
|
})
|
|
refused = socket.receive_json()
|
|
self.assertEqual(refused["code"], "CROP_TOO_LARGE")
|
|
socket.send_json({"type": "ping", "t": 1})
|
|
self.assertEqual(socket.receive_json(), {"type": "pong", "t": 1})
|
|
|
|
def test_a_crop_without_an_image_is_refused(self):
|
|
with self.client.websocket_connect("/v1/stream") as socket:
|
|
socket.receive_json() # hello
|
|
socket.send_json({"type": "crop", "request_id": "c1", "track_id": 1, "frame_id": 1})
|
|
refused = socket.receive_json()
|
|
self.assertEqual(refused["code"], "CROP_BAD_MESSAGE")
|
|
socket.send_json({"type": "ping", "t": 1})
|
|
self.assertEqual(socket.receive_json(), {"type": "pong", "t": 1})
|
|
|
|
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()
|
|
|
|
def test_lifespan_shutdown_skips_a_reserved_slot(self):
|
|
"""A reservation is None until setup finishes. Shutdown must not call it."""
|
|
client, engine = make_app()
|
|
client.app.state.sessions["reserved"] = None
|
|
try:
|
|
with client:
|
|
pass
|
|
finally:
|
|
engine.shutdown()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|