import socket import unittest from configurable_udp_capture import ( DEFAULT_PACKET_LAYOUT, ConfigurablePacketAssembler, ConfigurableUdpCapture, normalize_packet_layout, normalize_packet_schema, packet_layout_from_schema, packet_schema_from_layout, ) def mik_style_packet(flags, sequence, packet_number, value, data=b""): return bytes((0, flags, sequence, packet_number)) + int(value).to_bytes(4, "little") + data class ConfigurableUdpCaptureTests(unittest.TestCase): def test_visual_layout_calculates_offsets_and_arbitrary_field_sizes(self): layout = [ {"role": "skip", "size": 2, "label": "magic"}, {"role": "sequence", "size": 3, "label": "frame id"}, {"role": "flags", "size": 2, "label": "flags"}, {"role": "value", "size": 4, "label": "length"}, ] schema = packet_schema_from_layout(layout, {"byte_order": "big"}) self.assertEqual(schema["header_size"], 11) self.assertEqual((schema["sequence_offset"], schema["sequence_size"]), (2, 3)) self.assertEqual((schema["flags_offset"], schema["flags_size"]), (5, 2)) self.assertEqual((schema["value_offset"], schema["value_size"]), (7, 4)) self.assertEqual(schema["byte_order"], "big") def test_old_offset_schema_converts_to_visual_layout(self): layout = packet_layout_from_schema({ "header_size": 10, "flags_offset": 2, "flags_size": 1, "sequence_offset": 4, "sequence_size": 2, "packet_number_offset": -1, "value_offset": 6, "value_size": 4, }) self.assertEqual(sum(field["size"] for field in layout), 10) self.assertEqual([field["role"] for field in layout], ["skip", "flags", "skip", "sequence", "value"]) rebuilt = packet_schema_from_layout(layout) self.assertEqual(rebuilt["flags_offset"], 2) self.assertEqual(rebuilt["sequence_offset"], 4) self.assertEqual(rebuilt["value_offset"], 6) def test_empty_visual_layout_means_zero_byte_header(self): self.assertEqual(normalize_packet_layout([]), []) self.assertEqual(packet_schema_from_layout([], {"assembly": "datagram"})["header_size"], 0) self.assertEqual(sum(field["size"] for field in DEFAULT_PACKET_LAYOUT), 8) def test_named_visual_field_is_read_as_integer(self): schema = packet_schema_from_layout( [ {"role": "field", "size": 2, "label": "temperature"}, {"role": "skip", "size": 1, "label": "reserved"}, ], {"assembly": "datagram", "byte_order": "big"}, ) assembler = ConfigurablePacketAssembler(schema) self.assertEqual(assembler.push(b"\x01\x02\xffpayload"), b"payload") self.assertEqual(assembler.last_fields, {"temperature": 0x0102}) def test_schema_disables_fields_outside_header(self): schema = normalize_packet_schema({ "header_size": 2, "flags_offset": 1, "sequence_offset": 2, "packet_number_offset": 9, "value_offset": 4, "start_mask": "0x80", }) self.assertEqual(schema["flags_offset"], 1) self.assertEqual(schema["sequence_offset"], -1) self.assertEqual(schema["packet_number_offset"], -1) self.assertEqual(schema["value_offset"], -1) self.assertEqual(schema["start_mask"], 0x80) def test_fragmented_packet_schema_reassembles_raw_frame(self): frame_data = bytes(range(8)) cap = ConfigurableUdpCapture( host="127.0.0.1", port=0, width=4, height=2, encoding="gray8", ) sender = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) try: address = ("127.0.0.1", cap.port) sender.sendto(mik_style_packet(2, 7, 0, 8, frame_data[:3]), address) sender.sendto(mik_style_packet(0, 7, 1, 3, frame_data[3:6]), address) sender.sendto(mik_style_packet(1, 7, 2, 6, frame_data[6:]), address) ok, frame = cap.read() self.assertTrue(ok) self.assertEqual(frame.shape, (2, 4, 3)) self.assertEqual(int(frame[-1, -1, 0]), 7) finally: sender.close() cap.release() def test_datagram_schema_removes_header(self): frame_data = bytes(range(24)) cap = ConfigurableUdpCapture( host="127.0.0.1", port=0, width=4, height=2, encoding="bgr24", schema={"assembly": "datagram", "header_size": 2}, ) sender = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) try: sender.sendto(b"\xaa\x55" + frame_data, ("127.0.0.1", cap.port)) ok, frame = cap.read() self.assertTrue(ok) self.assertEqual(frame.shape, (2, 4, 3)) self.assertEqual(int(frame[0, 0, 0]), 0) finally: sender.close() cap.release() def test_custom_transport_can_decode_mik_video_array(self): width, height = 4, 2 mik_array = ( (0).to_bytes(4, "little") + width.to_bytes(2, "little") + height.to_bytes(2, "little") + bytes((1, 0, 0, 0)) + bytes(range(width * height)) ) cap = ConfigurableUdpCapture( host="127.0.0.1", port=0, schema={ "assembly": "datagram", "payload_format": "mik", "header_size": 0, }, ) sender = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) try: sender.sendto(mik_array, ("127.0.0.1", cap.port)) ok, frame = cap.read() self.assertTrue(ok) self.assertEqual(frame.shape, (height, width, 3)) finally: sender.close() cap.release() def test_stream_layout_reads_named_header_and_fixed_raw_frame(self): layout = [{"role": "field", "size": 2, "label": "camera_id"}] cap = ConfigurableUdpCapture( host="127.0.0.1", port=0, width=4, height=2, encoding="gray8", schema=packet_schema_from_layout(layout, {"assembly": "stream"}), ) sender = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) try: sender.sendto(b"\x2a\x00" + bytes(range(8)), ("127.0.0.1", cap.port)) ok, frame = cap.read() self.assertTrue(ok) self.assertEqual(frame.shape, (2, 4, 3)) self.assertEqual(cap.last_packet_fields, {"camera_id": 42}) finally: sender.close() cap.release() def test_bad_fragment_is_dropped_and_next_start_recovers(self): assembler = ConfigurablePacketAssembler() self.assertIsNone(assembler.push(mik_style_packet(2, 1, 0, 4, b"\x01\x02"))) self.assertIsNone(assembler.push(mik_style_packet(1, 1, 9, 2, b"\x03\x04"))) result = assembler.push(mik_style_packet(3, 2, 0, 4, b"\x05\x06\x07\x08")) self.assertEqual(result, b"\x05\x06\x07\x08") self.assertEqual(assembler.dropped_arrays, 1) def test_multibyte_big_endian_flags_and_short_packet_reset(self): assembler = ConfigurablePacketAssembler({ "header_size": 2, "byte_order": "big", "flags_offset": 0, "flags_size": 2, "start_mask": 0x8000, "end_mask": 0x4000, "sequence_offset": -1, "packet_number_offset": -1, "value_offset": -1, "value_mode": "unused", }) self.assertIsNone(assembler.push(b"\x80\x00first")) with self.assertRaises(ValueError): assembler.push(b"\x00") self.assertIsNone(assembler.push(b"\x40\x00ignored")) self.assertEqual(assembler.push(b"\xc0\x00next"), b"next") self.assertEqual(assembler.dropped_arrays, 1) if __name__ == "__main__": unittest.main()