mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 12:36:36 +00:00
[AZ-161] [AZ-162] [AZ-163] Add ONNX inference, NMS, annotation queue tests
Made-with: Cursor
This commit is contained in:
@@ -0,0 +1,76 @@
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import msgpack
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "annotation-queue"))
|
||||
from annotation_queue_dto import AnnotationBulkMessage, AnnotationMessage, AnnotationStatus, RoleEnum
|
||||
|
||||
|
||||
def _pack_created_message():
|
||||
ts = msgpack.Timestamp.from_datetime(datetime(2024, 1, 1, 12, 0, 0))
|
||||
detections_json = json.dumps(
|
||||
[{"an": "test", "cl": 0, "x": 0.5, "y": 0.5, "w": 0.1, "h": 0.1, "p": 0.9}]
|
||||
)
|
||||
data = [
|
||||
ts,
|
||||
"test-annotation",
|
||||
"media001.jpg",
|
||||
10000000,
|
||||
".jpg",
|
||||
detections_json,
|
||||
b"\xff\xd8\xff\xe0",
|
||||
20,
|
||||
"test@example.com",
|
||||
0,
|
||||
10,
|
||||
]
|
||||
return msgpack.packb(data)
|
||||
|
||||
|
||||
def test_bt_aqm_01_parse_created_annotation_message():
|
||||
msg = AnnotationMessage(_pack_created_message())
|
||||
assert msg.name == "test-annotation"
|
||||
assert len(msg.detections) == 1
|
||||
assert msg.detections[0].annotation_name == "test"
|
||||
assert msg.status == AnnotationStatus.Created
|
||||
assert msg.createdRole == RoleEnum.Validator
|
||||
assert msg.image == b"\xff\xd8\xff\xe0"
|
||||
|
||||
|
||||
def test_bt_aqm_02_parse_validated_bulk_message():
|
||||
ts = msgpack.Timestamp.from_datetime(datetime(2024, 1, 1, 12, 0, 0))
|
||||
packed = msgpack.packb([["n1", "n2"], 30, "admin@example.com", ts])
|
||||
msg = AnnotationBulkMessage(packed)
|
||||
assert msg.annotation_status == AnnotationStatus.Validated
|
||||
assert msg.annotation_names == ["n1", "n2"]
|
||||
|
||||
|
||||
def test_bt_aqm_03_parse_deleted_bulk_message():
|
||||
ts = msgpack.Timestamp.from_datetime(datetime(2024, 1, 1, 12, 0, 0))
|
||||
packed = msgpack.packb([["d1", "d2", "d3"], 40, "user@example.com", ts])
|
||||
msg = AnnotationBulkMessage(packed)
|
||||
assert msg.annotation_status == AnnotationStatus.Deleted
|
||||
assert msg.annotation_names == ["d1", "d2", "d3"]
|
||||
|
||||
|
||||
def test_bt_aqm_04_malformed_message_raises():
|
||||
with pytest.raises(Exception):
|
||||
AnnotationMessage(os.urandom(512))
|
||||
|
||||
|
||||
def _unpack_or_catch(raw):
|
||||
try:
|
||||
msgpack.unpackb(raw, strict_map_key=False)
|
||||
except Exception:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
@pytest.mark.resilience
|
||||
def test_rt_aqm_01_malformed_msgpack_bytes_handled():
|
||||
assert _unpack_or_catch(os.urandom(256))
|
||||
Reference in New Issue
Block a user