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.
208 lines
8.0 KiB
Python
208 lines
8.0 KiB
Python
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()
|