from __future__ import annotations import unittest from datetime import datetime, timezone from cold_display_guard.alarm_snapshots import ( capture_alert_snapshot, 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"]) if __name__ == "__main__": unittest.main()