316 lines
13 KiB
Python
316 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
from cold_display_guard import alarm_snapshots
|
|
from cold_display_guard.alarm_snapshots import (
|
|
capture_alert_snapshot,
|
|
fallback_label_text,
|
|
load_alarm_snapshot_settings,
|
|
upload_snapshot_bytes,
|
|
)
|
|
from cold_display_guard.vision import Frame
|
|
|
|
|
|
UTC = timezone.utc
|
|
|
|
|
|
class AlarmSnapshotTests(unittest.TestCase):
|
|
def test_load_alarm_snapshot_settings_from_config(self) -> None:
|
|
settings = load_alarm_snapshot_settings(
|
|
{
|
|
"alarm_snapshot_upload": {
|
|
"enabled": True,
|
|
"service_url": "https://ota.zhengxinshipin.com",
|
|
"secret": "change-me-in-production",
|
|
"object_key_prefix": "alarms/cold-display",
|
|
"connect_timeout_seconds": 4,
|
|
"read_timeout_seconds": 9,
|
|
"encode_timeout_seconds": 7,
|
|
}
|
|
}
|
|
)
|
|
|
|
self.assertTrue(settings.enabled)
|
|
self.assertEqual(settings.service_url, "https://ota.zhengxinshipin.com")
|
|
self.assertEqual(settings.secret, "change-me-in-production")
|
|
self.assertEqual(settings.object_key_prefix, "alarms/cold-display")
|
|
self.assertEqual(settings.connect_timeout_seconds, 4)
|
|
self.assertEqual(settings.read_timeout_seconds, 9)
|
|
self.assertEqual(settings.encode_timeout_seconds, 7)
|
|
|
|
def test_upload_snapshot_bytes_uses_documented_chunk_upload_flow(self) -> None:
|
|
json_calls: list[tuple[str, dict[str, object]]] = []
|
|
chunk_calls: list[tuple[str, dict[str, str], bytes, dict[str, str]]] = []
|
|
|
|
def fake_post_json(url: str, payload: dict[str, object], timeout: tuple[float, float]) -> tuple[int, dict[str, object]]:
|
|
json_calls.append((url, payload))
|
|
if url.endswith("/token/generate"):
|
|
return 200, {"token": "token-1", "expires_at": 1770003600}
|
|
if url.endswith("/upload/init"):
|
|
return 200, {"upload_id": "upload-1"}
|
|
if url.endswith("/upload/complete"):
|
|
return 200, {"upload_id": "upload-1", "object_key": "uploads/alarms/a.jpg", "file_size": 3, "file_md5": "900150983cd24fb0d6963f7d28e17f72"}
|
|
raise AssertionError(url)
|
|
|
|
def fake_post_multipart(
|
|
url: str,
|
|
fields: dict[str, str],
|
|
file_field: str,
|
|
file_name: str,
|
|
file_bytes: bytes,
|
|
timeout: tuple[float, float],
|
|
) -> tuple[int, dict[str, object]]:
|
|
chunk_calls.append((url, fields, file_bytes, {"file_field": file_field, "file_name": file_name}))
|
|
return 200, {"upload_id": "upload-1", "index": 0, "size": len(file_bytes), "received_chunks": 1, "total_chunks": 1}
|
|
|
|
result = upload_snapshot_bytes(
|
|
b"abc",
|
|
file_name="alarm.jpg",
|
|
object_key_hint="alarms/a.jpg",
|
|
settings=load_alarm_snapshot_settings({}),
|
|
post_json_request=fake_post_json,
|
|
post_multipart_request=fake_post_multipart,
|
|
)
|
|
|
|
self.assertEqual(result["status"], "uploaded")
|
|
self.assertEqual(result["object_key"], "uploads/alarms/a.jpg")
|
|
self.assertEqual(json_calls[0][0], "https://ota.zhengxinshipin.com/token/generate")
|
|
self.assertEqual(json_calls[1][0], "https://ota.zhengxinshipin.com/upload/init")
|
|
self.assertEqual(json_calls[2][0], "https://ota.zhengxinshipin.com/upload/complete")
|
|
self.assertIn("token=token-1", chunk_calls[0][0])
|
|
self.assertIn("upload_id=upload-1", chunk_calls[0][0])
|
|
self.assertEqual(chunk_calls[0][1]["chunk_md5"], "900150983cd24fb0d6963f7d28e17f72")
|
|
self.assertEqual(chunk_calls[0][3]["file_field"], "chunk")
|
|
|
|
def test_capture_alert_snapshot_skips_non_alert_events(self) -> None:
|
|
result = capture_alert_snapshot(
|
|
Frame(width=1, height=1, rgb=b"\x00\x00\x00"),
|
|
[{"event": "batch_started", "severity": "info", "batch_id": "batch_1"}],
|
|
{},
|
|
now=datetime(2026, 6, 9, 9, 0, tzinfo=UTC),
|
|
)
|
|
|
|
self.assertIsNone(result)
|
|
|
|
def test_capture_alert_snapshot_uploads_current_frame_for_alert_events(self) -> None:
|
|
encode_calls: list[Frame] = []
|
|
upload_calls: list[tuple[bytes, str, str]] = []
|
|
|
|
def fake_encode(frame: Frame, timeout_seconds: float) -> bytes:
|
|
encode_calls.append(frame)
|
|
return b"jpeg-bytes"
|
|
|
|
def fake_upload(
|
|
image_bytes: bytes,
|
|
*,
|
|
file_name: str,
|
|
object_key_hint: str,
|
|
settings,
|
|
post_json_request=None,
|
|
post_multipart_request=None,
|
|
) -> dict[str, object]:
|
|
upload_calls.append((image_bytes, file_name, object_key_hint))
|
|
return {"status": "uploaded", "object_key": "uploads/alarms/test.jpg", "file_name": file_name}
|
|
|
|
result = capture_alert_snapshot(
|
|
Frame(width=1, height=1, rgb=b"\x01\x02\x03"),
|
|
[{"event": "time_alarm", "severity": "alarm", "batch_id": "batch_1", "camera_id": "cam_1", "zone_id": "1", "ts": "2026-06-09T09:00:00+00:00"}],
|
|
{"alarm_snapshot_upload": {"enabled": True, "object_key_prefix": "alarms/cold-display"}},
|
|
now=datetime(2026, 6, 9, 9, 0, tzinfo=UTC),
|
|
jpeg_encoder=fake_encode,
|
|
uploader=fake_upload,
|
|
)
|
|
|
|
self.assertEqual(len(encode_calls), 1)
|
|
self.assertEqual(upload_calls[0][0], b"jpeg-bytes")
|
|
self.assertEqual(result["status"], "uploaded")
|
|
self.assertEqual(result["object_key"], "uploads/alarms/test.jpg")
|
|
self.assertEqual(result["batch_ids"], ["batch_1"])
|
|
|
|
def test_calibration_overlay_draws_zones_and_trash_roi_without_mutating_source(self) -> None:
|
|
apply_overlay = getattr(alarm_snapshots, "apply_calibration_overlay", None)
|
|
self.assertTrue(callable(apply_overlay), "apply_calibration_overlay should be available")
|
|
frame = Frame(width=5, height=5, rgb=b"\x00\x00\x00" * 25)
|
|
|
|
annotated = apply_overlay(
|
|
frame,
|
|
{
|
|
"zones": [
|
|
{
|
|
"id": "1",
|
|
"label": "区域 1",
|
|
"polygon": [[0.2, 0.2], [0.8, 0.2], [0.8, 0.8], [0.2, 0.8]],
|
|
}
|
|
],
|
|
"trash": {"roi": [[0.0, 0.0], [0.4, 0.0], [0.0, 0.4]]},
|
|
},
|
|
)
|
|
|
|
self.assertEqual(frame.rgb, b"\x00\x00\x00" * 25)
|
|
self.assertNotEqual(annotated.rgb, frame.rgb)
|
|
self.assertNotEqual(annotated.pixel(1, 1), (0, 0, 0))
|
|
self.assertNotEqual(annotated.pixel(2, 2), (0, 0, 0))
|
|
self.assertNotEqual(annotated.pixel(0, 0), (0, 0, 0))
|
|
|
|
def test_calibration_overlay_uses_distinct_zone_colors_and_draws_labels(self) -> None:
|
|
frame = Frame(width=40, height=20, rgb=b"\x00\x00\x00" * 800)
|
|
|
|
annotated = alarm_snapshots.apply_calibration_overlay(
|
|
frame,
|
|
{
|
|
"zones": [
|
|
{
|
|
"id": "1",
|
|
"label": "区域 1",
|
|
"polygon": [[0.05, 0.10], [0.40, 0.10], [0.40, 0.90], [0.05, 0.90]],
|
|
},
|
|
{
|
|
"id": "2",
|
|
"label": "区域 2",
|
|
"polygon": [[0.55, 0.10], [0.90, 0.10], [0.90, 0.90], [0.55, 0.90]],
|
|
},
|
|
]
|
|
},
|
|
)
|
|
|
|
self.assertNotEqual(annotated.pixel(10, 15), annotated.pixel(30, 15))
|
|
label_pixels = [annotated.pixel(x, y) for y in range(2, 10) for x in range(2, 18)]
|
|
self.assertTrue(any(max(pixel) >= 220 for pixel in label_pixels), "expected bright label text pixels")
|
|
|
|
def test_chinese_label_fallback_uses_readable_ascii_when_font_renderer_is_unavailable(self) -> None:
|
|
self.assertEqual(fallback_label_text("区域 1"), "R1")
|
|
self.assertEqual(fallback_label_text("区域 12"), "R12")
|
|
self.assertEqual(fallback_label_text("垃圾区"), "TRASH")
|
|
|
|
def test_docker_image_installs_cjk_fonts_for_alarm_snapshot_labels(self) -> None:
|
|
dockerfile = (Path(__file__).resolve().parents[1] / "Dockerfile").read_text(encoding="utf-8")
|
|
self.assertIn("fonts-noto-cjk", dockerfile)
|
|
|
|
def test_drawtext_label_style_stays_small_and_translucent(self) -> None:
|
|
filter_text = alarm_snapshots.build_drawtext_filter(
|
|
[
|
|
alarm_snapshots.OverlayLabel(
|
|
text="区域 1",
|
|
fallback_text="R1",
|
|
x=10,
|
|
y=20,
|
|
accent_rgb=(255, 196, 0),
|
|
)
|
|
],
|
|
Path("/tmp/NotoSansCJK-Regular.ttc"),
|
|
height=360,
|
|
)
|
|
|
|
self.assertIn("fontsize=13", filter_text)
|
|
self.assertIn("boxcolor=black@0.34", filter_text)
|
|
self.assertIn("boxborderw=2", filter_text)
|
|
|
|
def test_capture_alert_snapshot_encodes_frame_with_configured_calibration_overlay(self) -> None:
|
|
encoded_frames: list[Frame] = []
|
|
|
|
def fake_encode(frame: Frame, timeout_seconds: float) -> bytes:
|
|
encoded_frames.append(frame)
|
|
return b"jpeg-bytes"
|
|
|
|
def fake_upload(
|
|
image_bytes: bytes,
|
|
*,
|
|
file_name: str,
|
|
object_key_hint: str,
|
|
settings,
|
|
post_json_request=None,
|
|
post_multipart_request=None,
|
|
) -> dict[str, object]:
|
|
return {"status": "uploaded", "object_key": "uploads/alarms/overlay.jpg", "file_name": file_name}
|
|
|
|
source_frame = Frame(width=5, height=5, rgb=b"\x00\x00\x00" * 25)
|
|
result = capture_alert_snapshot(
|
|
source_frame,
|
|
[{"event": "time_alarm", "severity": "alarm", "batch_id": "batch_1", "camera_id": "cam_1", "zone_id": "1", "ts": "2026-06-09T09:00:00+00:00"}],
|
|
{
|
|
"alarm_snapshot_upload": {"enabled": True},
|
|
"zones": [
|
|
{
|
|
"id": "1",
|
|
"label": "区域 1",
|
|
"polygon": [[0.2, 0.2], [0.8, 0.2], [0.8, 0.8], [0.2, 0.8]],
|
|
}
|
|
],
|
|
"trash": {"roi": [[0.0, 0.0], [0.4, 0.0], [0.0, 0.4]]},
|
|
},
|
|
now=datetime(2026, 6, 9, 9, 0, tzinfo=UTC),
|
|
jpeg_encoder=fake_encode,
|
|
uploader=fake_upload,
|
|
)
|
|
|
|
self.assertEqual(result["status"], "uploaded")
|
|
self.assertEqual(source_frame.rgb, b"\x00\x00\x00" * 25)
|
|
self.assertEqual(len(encoded_frames), 1)
|
|
self.assertNotEqual(encoded_frames[0].rgb, source_frame.rgb)
|
|
self.assertNotEqual(encoded_frames[0].pixel(1, 1), (0, 0, 0))
|
|
|
|
def test_capture_alert_snapshot_only_draws_alert_event_zones(self) -> None:
|
|
encoded_frames: list[Frame] = []
|
|
|
|
def fake_encode(frame: Frame, timeout_seconds: float) -> bytes:
|
|
encoded_frames.append(frame)
|
|
return b"jpeg-bytes"
|
|
|
|
def fake_upload(
|
|
image_bytes: bytes,
|
|
*,
|
|
file_name: str,
|
|
object_key_hint: str,
|
|
settings,
|
|
post_json_request=None,
|
|
post_multipart_request=None,
|
|
) -> dict[str, object]:
|
|
return {"status": "uploaded", "object_key": "uploads/alarms/zone-only.jpg", "file_name": file_name}
|
|
|
|
source_frame = Frame(width=30, height=20, rgb=b"\x00\x00\x00" * 600)
|
|
result = capture_alert_snapshot(
|
|
source_frame,
|
|
[
|
|
{
|
|
"event": "time_alarm",
|
|
"severity": "alarm",
|
|
"batch_id": "batch_1",
|
|
"camera_id": "cam_1",
|
|
"zone_id": "1",
|
|
"ts": "2026-06-09T09:00:00+00:00",
|
|
}
|
|
],
|
|
{
|
|
"alarm_snapshot_upload": {"enabled": True},
|
|
"zones": [
|
|
{
|
|
"id": "1",
|
|
"label": "区域 1",
|
|
"polygon": [[0.00, 0.00], [0.45, 0.00], [0.45, 1.00], [0.00, 1.00]],
|
|
},
|
|
{
|
|
"id": "2",
|
|
"label": "区域 2",
|
|
"polygon": [[0.55, 0.00], [1.00, 0.00], [1.00, 1.00], [0.55, 1.00]],
|
|
},
|
|
],
|
|
"trash": {"roi": [[0.45, 0.50], [0.55, 0.50], [0.55, 1.00], [0.45, 1.00]]},
|
|
},
|
|
now=datetime(2026, 6, 9, 9, 0, tzinfo=UTC),
|
|
jpeg_encoder=fake_encode,
|
|
uploader=fake_upload,
|
|
)
|
|
|
|
self.assertEqual(result["status"], "uploaded")
|
|
self.assertEqual(len(encoded_frames), 1)
|
|
self.assertNotEqual(encoded_frames[0].pixel(5, 10), (0, 0, 0))
|
|
self.assertEqual(encoded_frames[0].pixel(25, 10), (0, 0, 0))
|
|
self.assertEqual(encoded_frames[0].pixel(15, 15), (0, 0, 0))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|