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_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)) if __name__ == "__main__": unittest.main()