NODEDC_MISSION_CORE/tests/test_stream_summary.py

196 lines
6.0 KiB
Python

from __future__ import annotations
import hashlib
import json
import stat
import struct
from pathlib import Path
import lz4.block
import pytest
from typer.testing import CliRunner
from k1link.analyze.stream_summary import summarize_mqtt_streams
from k1link.cli import app
from k1link.mqtt.capture import FRAME_HEADER, RAW_MAGIC, CaptureFormatError
runner = CliRunner()
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 _fixed64(number: int, value: float) -> bytes:
return _key(number, 1) + struct.pack("<d", value)
def _header(*, scaler: int) -> bytes:
return b"".join(
(
_uint(1, 7),
_sint(2, 123456),
_sint(3, scaler),
_bytes(4, b"device-super-secret"),
_bytes(5, b"session-super-secret"),
_bytes(6, b"openapi-super-secret"),
)
)
def _pcl_payload(*, scaler: int, point_count: int) -> bytes:
points = []
for index in range(point_count):
point = (
_sint(1, 1000 + index)
+ _sint(2, -2000 - index)
+ _sint(3, 500 + index)
+ _uint(4, index)
)
points.append(_bytes(2, point))
report = _bytes(1, _header(scaler=scaler)) + b"".join(points)
compressed = lz4.block.compress(report, store_size=False)
return _uint(3, len(report)) + _bytes(4, compressed)
def _pose_payload(position_xyz: tuple[float, float, float]) -> bytes:
position = b"".join(
_fixed64(field_number, value)
for field_number, value in enumerate(position_xyz, start=1)
)
orientation = (
_fixed64(1, 0.0)
+ _fixed64(2, 0.0)
+ _fixed64(3, 0.0)
+ _fixed64(4, 1.0)
)
pose = _bytes(1, position) + _bytes(2, orientation)
stamped = _sint(1, 987654321) + _bytes(2, pose)
return _bytes(1, _header(scaler=1000)) + _bytes(2, stamped)
def _write_capture(path: Path, frames: list[tuple[str, bytes]]) -> bytes:
raw = bytearray(RAW_MAGIC)
for topic, payload in frames:
topic_bytes = topic.encode()
raw.extend(FRAME_HEADER.pack(len(topic_bytes), len(payload)))
raw.extend(topic_bytes)
raw.extend(payload)
value = bytes(raw)
path.write_bytes(value)
return value
def test_summary_is_bounded_aggregate_only_and_hashes_capture(tmp_path: Path) -> None:
frames = [
("lixel/application/report/lio_pcl", _pcl_payload(scaler=1000, point_count=2)),
("lixel/application/report/lio_pcl", _pcl_payload(scaler=2000, point_count=1)),
("lixel/application/report/lio_pcl", b"bad-protobuf"),
("lixel/application/report/lio_pose", _pose_payload((1.0, 2.0, 3.0))),
("lixel/application/report/lio_pose", _pose_payload((4.0, 6.0, 3.0))),
("lixel/application/report/lio_pose", b"bad-protobuf"),
("device-super-secret/session-super-secret", b"openapi-super-secret"),
]
capture = tmp_path / "mqtt.raw.k1mqtt"
raw = _write_capture(capture, frames)
summary = summarize_mqtt_streams(capture)
assert summary["source"] == {
"bytes": len(raw),
"sha256": hashlib.sha256(raw).hexdigest(),
}
assert summary["frames"] == {
"count": 7,
"payload_bytes": sum(len(payload) for _, payload in frames),
"encoded_frame_bytes": len(raw) - len(RAW_MAGIC),
"other_count": 1,
"other_payload_bytes": len(b"openapi-super-secret"),
}
assert summary["decoding"] == {"attempted": 6, "successes": 4, "errors": 2}
assert summary["point_cloud"]["frame_count"] == 3
assert summary["point_cloud"]["decode_successes"] == 2
assert summary["point_cloud"]["decode_errors"] == 1
assert summary["point_cloud"]["points"] == {
"total": 3,
"per_frame": {"min": 1, "max": 2},
}
assert summary["point_cloud"]["scalers"] == {
"min": 1000,
"max": 2000,
"constant": False,
}
assert summary["pose"]["frame_count"] == 3
assert summary["pose"]["decode_successes"] == 2
assert summary["pose"]["decode_errors"] == 1
assert summary["pose"]["first_to_last_displacement_meters"] == 5.0
serialized = json.dumps(summary)
for secret in (
"device-super-secret",
"session-super-secret",
"openapi-super-secret",
"bad-protobuf",
):
assert secret not in serialized
assert "position_xyz" not in serialized
assert "points" in summary["point_cloud"]
def test_summary_rejects_a_frame_above_the_operator_limit(tmp_path: Path) -> None:
capture = tmp_path / "mqtt.raw.k1mqtt"
_write_capture(capture, [("other", b"12345")])
with pytest.raises(CaptureFormatError, match="exceeds 4"):
summarize_mqtt_streams(capture, max_payload_bytes=4)
def test_mqtt_streams_cli_writes_atomic_private_summary(tmp_path: Path) -> None:
capture = tmp_path / "mqtt.raw.k1mqtt"
raw = _write_capture(capture, [])
out = tmp_path / "captures" / "summary.json"
result = runner.invoke(
app,
[
"analyze",
"mqtt-streams",
"--capture",
str(capture),
"--out",
str(out),
"--max-payload-bytes",
"1024",
],
)
assert result.exit_code == 0
saved = json.loads(out.read_text())
assert saved["source"]["sha256"] == hashlib.sha256(raw).hexdigest()
assert saved["limits"]["max_payload_bytes"] == 1024
assert saved["frames"]["count"] == 0
assert "aggregate-only" in result.stdout
assert stat.S_IMODE(out.stat().st_mode) == 0o600