160 lines
5.4 KiB
Python
160 lines
5.4 KiB
Python
from __future__ import annotations
|
|
|
|
import math
|
|
import struct
|
|
|
|
import lz4.block
|
|
import pytest
|
|
|
|
from k1link.device_plugins.xgrids_k1.protocol.protobuf_wire import (
|
|
ProtobufWireError,
|
|
decode_zigzag64,
|
|
iter_fields,
|
|
)
|
|
from k1link.device_plugins.xgrids_k1.protocol.streams import (
|
|
DecodeLimits,
|
|
StreamDecodeError,
|
|
UnsupportedCompressionError,
|
|
decode_legacy_pointcloud,
|
|
decode_legacy_pose,
|
|
decode_lio_pcl,
|
|
decode_lio_pose,
|
|
decode_pre_path_array,
|
|
)
|
|
|
|
|
|
def _varint(value: int) -> bytes:
|
|
encoded = bytearray()
|
|
while value > 0x7F:
|
|
encoded.append((value & 0x7F) | 0x80)
|
|
value >>= 7
|
|
encoded.append(value)
|
|
return bytes(encoded)
|
|
|
|
|
|
def _key(number: int, wire_type: int) -> bytes:
|
|
return _varint((number << 3) | wire_type)
|
|
|
|
|
|
def _uint(number: int, value: int) -> bytes:
|
|
return _key(number, 0) + _varint(value)
|
|
|
|
|
|
def _sint(number: int, value: int) -> bytes:
|
|
zigzag = (value << 1) ^ (value >> 63)
|
|
return _uint(number, zigzag & 0xFFFFFFFFFFFFFFFF)
|
|
|
|
|
|
def _bytes(number: int, value: bytes) -> bytes:
|
|
return _key(number, 2) + _varint(len(value)) + value
|
|
|
|
|
|
def _fixed32(number: int, value: float) -> bytes:
|
|
return _key(number, 5) + struct.pack("<f", value)
|
|
|
|
|
|
def _fixed64(number: int, value: float) -> bytes:
|
|
return _key(number, 1) + struct.pack("<d", value)
|
|
|
|
|
|
def _header(*, scaler: int = 1000) -> bytes:
|
|
return b"".join(
|
|
(
|
|
_uint(1, 7),
|
|
_sint(2, 123456),
|
|
_sint(3, scaler),
|
|
_bytes(4, b"device-redacted"),
|
|
_bytes(5, b"session-redacted"),
|
|
)
|
|
)
|
|
|
|
|
|
def _pcl_payload(*, compression: int = 0, scaler: int = 1000) -> bytes:
|
|
point_1 = _sint(1, 1000) + _sint(2, -2000) + _sint(3, 500) + _uint(4, 0x11223344)
|
|
point_2 = _sint(1, -250) + _sint(2, 0) + _sint(3, 4000) + _uint(4, 0xAABBCC09)
|
|
report = _bytes(1, _header(scaler=scaler)) + _bytes(2, point_1) + _bytes(2, point_2)
|
|
compressed = lz4.block.compress(report, store_size=False)
|
|
fields = []
|
|
if compression:
|
|
fields.append(_uint(2, compression))
|
|
fields.extend((_uint(3, len(report)), _bytes(4, compressed)))
|
|
return b"".join(fields)
|
|
|
|
|
|
def test_protobuf_wire_zigzag_and_bounds() -> None:
|
|
assert decode_zigzag64(0) == 0
|
|
assert decode_zigzag64(1) == -1
|
|
assert decode_zigzag64(2) == 1
|
|
with pytest.raises(ProtobufWireError, match="truncated"):
|
|
list(iter_fields(b"\x0a\x02\x01"))
|
|
with pytest.raises(ProtobufWireError, match="unsupported"):
|
|
list(iter_fields(b"\x0b"))
|
|
|
|
|
|
def test_decode_lio_pcl_raw_lz4() -> None:
|
|
frame = decode_lio_pcl(_pcl_payload())
|
|
assert frame.header.seq == 7
|
|
assert frame.header.stamp == 123456
|
|
assert frame.header.scaler == 1000
|
|
assert len(frame.points) == 2
|
|
assert frame.points[0].scaled_xyz(frame.header.scaler) == (1.0, -2.0, 0.5)
|
|
assert frame.points[0].rgbi == 0x11223344
|
|
assert frame.points[0].intensity == 0x44
|
|
assert frame.points[1].scaled_xyz(frame.header.scaler) == (-0.25, 0.0, 4.0)
|
|
assert frame.points[1].intensity == 9
|
|
|
|
|
|
def test_decode_lio_pcl_rejects_unverified_or_unsafe_frames() -> None:
|
|
with pytest.raises(UnsupportedCompressionError, match="enum 1"):
|
|
decode_lio_pcl(_pcl_payload(compression=1))
|
|
with pytest.raises(StreamDecodeError, match="scaler is zero"):
|
|
decode_lio_pcl(_pcl_payload(scaler=0))
|
|
with pytest.raises(StreamDecodeError, match="exceeds 1 points"):
|
|
decode_lio_pcl(_pcl_payload(), DecodeLimits(max_points_per_frame=1))
|
|
with pytest.raises(StreamDecodeError, match="MQTT payload exceeds"):
|
|
decode_lio_pcl(_pcl_payload(), DecodeLimits(max_mqtt_payload_bytes=4))
|
|
|
|
|
|
def test_decode_lio_pose() -> None:
|
|
position = _fixed64(1, 1.25) + _fixed64(2, -2.5) + _fixed64(3, 3.75)
|
|
orientation = _fixed64(1, 0.1) + _fixed64(2, 0.2) + _fixed64(3, 0.3) + _fixed64(4, 0.9)
|
|
pose = _bytes(1, position) + _bytes(2, orientation)
|
|
stamped = _sint(1, 987654321) + _bytes(2, pose)
|
|
payload = _bytes(1, _header()) + _bytes(2, stamped) + _fixed32(3, 12.5) + _fixed32(4, 0.001)
|
|
|
|
frame = decode_lio_pose(payload)
|
|
assert frame.pose_stamp == 987654321
|
|
assert frame.position_xyz == (1.25, -2.5, 3.75)
|
|
assert frame.orientation_xyzw == pytest.approx((0.1, 0.2, 0.3, 0.9))
|
|
assert frame.distance == 12.5
|
|
assert frame.pose_accuracy == pytest.approx(0.001)
|
|
|
|
|
|
def test_decode_lio_pose_rejects_nonfinite_float() -> None:
|
|
payload = _bytes(1, _header()) + _fixed32(3, math.nan)
|
|
with pytest.raises(StreamDecodeError, match="not finite"):
|
|
decode_lio_pose(payload)
|
|
|
|
|
|
def test_decode_legacy_pointcloud() -> None:
|
|
envelope = struct.pack("<III", 16, 123, 456)
|
|
body = struct.pack("<fffBBBB", 1.0, -2.0, 3.0, 10, 20, 30, 40)
|
|
frame = decode_legacy_pointcloud(envelope + body)
|
|
assert frame.stride == 16
|
|
assert frame.envelope == envelope
|
|
assert frame.points[0] == (1.0, -2.0, 3.0, 10, 20, 30, 40)
|
|
|
|
|
|
def test_decode_legacy_pose_and_matrix() -> None:
|
|
payload = struct.pack("<ffffffff", 1.0, 2.0, 3.0, 99.0, 0.9, 0.1, 0.2, 0.3)
|
|
frame = decode_legacy_pose(payload + b"tail")
|
|
assert frame.position_xyz == (1.0, 2.0, 3.0)
|
|
assert frame.orientation_xyzw == pytest.approx((0.1, 0.2, 0.3, 0.9))
|
|
assert frame.skipped_offset_12 == struct.pack("<f", 99.0)
|
|
assert frame.unknown_tail == b"tail"
|
|
|
|
matrix = tuple(float(index) for index in range(16))
|
|
assert decode_pre_path_array(struct.pack("<16d", *matrix)) == matrix
|
|
with pytest.raises(StreamDecodeError, match="exactly 128"):
|
|
decode_pre_path_array(b"short")
|