You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

75 lines
2.8 KiB
Python

import tempfile
import unittest
from pathlib import Path
from udp_probe import analyze_udp_records, save_udp_records
def mik_payload(flags, sequence, packet_number, value, data=b""):
return bytes((0, flags, sequence, packet_number)) + value.to_bytes(4, "little") + data
def mik_frame(width=4, height=2):
pixels = bytes(range(width * height))
image_header = width.to_bytes(2, "little") + height.to_bytes(2, "little") + bytes((1, 0, 0, 0))
return (0).to_bytes(4, "little") + image_header + pixels
def record(index, payload, source=("192.168.0.10", 40000)):
return {
"timestamp_ns": 1_000_000_000 + index * 1_000_000,
"address": source,
"payload": payload,
}
class UdpProbeTests(unittest.TestCase):
def test_detects_complete_mik_video_array(self):
frame = mik_frame()
split = 9
payloads = [
mik_payload(2, 7, 0, len(frame)),
mik_payload(0, 7, 1, 0, frame[:split]),
mik_payload(0, 7, 2, split, frame[split:]),
mik_payload(1, 7, 3, len(frame)),
]
result = analyze_udp_records([record(index, payload) for index, payload in enumerate(payloads)])
self.assertEqual(result["detected"]["kind"], "mik_video")
self.assertEqual(result["detected"]["confidence"], 100)
self.assertEqual(result["detected"]["frame"]["width"], 4)
self.assertEqual(result["detected"]["recommended"]["source_mode"], "udp_mik_live")
def test_detects_selected_raw_frame_with_separator_packet(self):
payloads = [b"\x01\x02\x03\x04", b"\x05\x06\x07\x08", b"\xff"]
result = analyze_udp_records(
[record(index, payload) for index, payload in enumerate(payloads)],
width=4,
height=2,
separator=255,
frame_encoding="gray8",
)
self.assertEqual(result["detected"]["kind"], "raw_delimited")
self.assertEqual(result["detected"]["confidence"], 99)
self.assertEqual(result["detected"]["recommended"]["frame_encoding"], "gray8")
def test_saved_dump_preserves_every_payload_byte(self):
records = [record(0, b"\x00\x01"), record(1, b"\xfe\xff")]
with tempfile.TemporaryDirectory() as tmp:
dump_path, report_path, report = save_udp_records(records, 40404, Path(tmp))
expected = (
(40404).to_bytes(2, "little") + (2).to_bytes(2, "little") + b"\x00\x01"
+ (40404).to_bytes(2, "little") + (2).to_bytes(2, "little") + b"\xfe\xff"
)
self.assertEqual(dump_path.read_bytes(), expected)
self.assertTrue(report_path.is_file())
self.assertEqual(report["dump_size"], len(expected))
self.assertEqual(len(report["dump_sha256"]), 64)
if __name__ == "__main__":
unittest.main()