NODEDC_MISSION_CORE/tests/test_rerun_bridge.py

from __future__ import annotations

import json
import socket
import struct
import threading
import time
from pathlib import Path

import numpy as np
import pytest

from k1link.mqtt.capture import FRAME_HEADER, RAW_MAGIC
from k1link.viewer.messages import StreamMessage
from k1link.viewer.rerun_bridge import RerunBridge, RerunSceneSettings, _point_colors
from k1link.viewer.runtime import VisualizationRuntime


class FakeRecording:
    def __init__(self) -> None:
        self.logs: list[tuple[str, object, bool]] = []
        self.times: list[tuple[str, dict[str, object]]] = []
        self.blueprints: list[object] = []
        self.disconnected = False

    def serve_grpc(self, **_: object) -> str:
        return "rerun+http://127.0.0.1:9876/proxy"

    def log(self, path: str, entity: object, *, static: bool = False) -> None:
        self.logs.append((path, entity, static))

    def set_time(self, timeline: str, **value: object) -> None:
        self.times.append((timeline, value))

    def send_blueprint(self, blueprint: object, **_: object) -> None:
        self.blueprints.append(blueprint)

    def disconnect(self) -> None:
        self.disconnected = True

    def flush(self, **_: object) -> None:
        return


def _message(topic: str, payload: bytes, *, sequence: int = 7) -> StreamMessage:
    return StreamMessage(
        sequence=sequence,
        topic=topic,
        payload=payload,
        received_at_epoch_ns=1_784_124_315_186_225_000,
        received_monotonic_ns=None,
        source="test",
    )


def test_legacy_points_and_pose_are_logged_to_rerun() -> None:
    recording = FakeRecording()
    bridge = RerunBridge(recording_factory=lambda _: recording)  # type: ignore[arg-type]

    point_payload = struct.pack("<III", 16, 0, 0) + struct.pack(
        "<fffBBBB", 1.0, -2.0, 3.0, 10, 20, 30, 40
    )
    pose_payload = struct.pack("<ffffffff", 1.0, 2.0, 3.0, 99.0, 0.9, 0.1, 0.2, 0.3)
    bridge.process(_message("RealtimePointcloud", point_payload))
    bridge.process(_message("RealtimePath", pose_payload, sequence=8))

    paths = [path for path, _, _ in recording.logs]
    snapshot = bridge.metrics.snapshot()
    assert bridge.grpc_url == "rerun+http://127.0.0.1:9876/proxy"
    assert "/world/points" in paths
    assert "/world/sensor_pose" in paths
    assert "/world/trajectory" in paths
    assert snapshot["pcl_frames"] == 1
    assert snapshot["pose_frames"] == 1
    assert snapshot["last_point_count"] == 1
    assert {timeline for timeline, _ in recording.times} == {
        "capture_time",
        "message_sequence",
        "stream_time",
    }

    bridge.close()
    assert recording.disconnected is True


def test_bad_frame_is_counted_without_publishing() -> None:
    recording = FakeRecording()
    bridge = RerunBridge(recording_factory=lambda _: recording)  # type: ignore[arg-type]

    bridge.process(_message("RealtimePointcloud", b"short"))

    snapshot = bridge.metrics.snapshot()
    assert snapshot["messages_received"] == 1
    assert snapshot["decode_errors"] == 1
    assert not any(path == "/world/points" for path, _, _ in recording.logs)


def test_palettes_are_deterministic_and_custom_color_is_exact() -> None:
    positions = np.asarray([[0, 0, 0], [0, 0, 10]], dtype=np.float32)
    intensities = np.asarray([0, 255], dtype=np.uint8)

    height = _point_colors(
        positions,
        intensities,
        None,
        RerunSceneSettings(color_mode="height", palette="viridis"),
    )
    custom = _point_colors(
        positions,
        intensities,
        None,
        RerunSceneSettings(color_mode="class", custom_color="#102030"),
    )

    assert height.shape == (2, 3)
    assert height.dtype == np.uint8
    assert height[0].tolist() == [68, 1, 84]
    assert height[1].tolist() == [253, 231, 37]
    assert custom.tolist() == [[16, 32, 48], [16, 32, 48]]


def test_real_grpc_server_releases_its_port() -> None:
    probe = socket.socket()
    probe.bind(("127.0.0.1", 0))
    port = int(probe.getsockname()[1])
    probe.close()

    bridge = RerunBridge(
        grpc_port=port,
        cors_allow_origin=("http://127.0.0.1:8000",),
    )
    assert bridge.grpc_url == f"rerun+http://127.0.0.1:{port}/proxy"
    bridge.close()

    deadline = time.monotonic() + 2.0
    while True:
        available = socket.socket()
        try:
            available.bind(("127.0.0.1", port))
            break
        except OSError:
            if time.monotonic() >= deadline:
                raise
            time.sleep(0.02)
        finally:
            available.close()


def test_runtime_exposes_rerun_url_and_stops_cleanly(tmp_path: Path) -> None:
    capture = tmp_path / "mqtt.raw.k1mqtt"
    point_topic = "RealtimePointcloud"
    pose_topic = "RealtimePath"
    point_payload = struct.pack("<III", 16, 0, 0) + struct.pack(
        "<fffBBBB", 1.0, -2.0, 3.0, 10, 20, 30, 40
    )
    pose_payload = struct.pack("<ffffffff", 1.0, 2.0, 3.0, 99.0, 0.9, 0.1, 0.2, 0.3)
    frames = bytearray(RAW_MAGIC)
    for topic, payload in ((point_topic, point_payload), (pose_topic, pose_payload)):
        topic_raw = topic.encode()
        frames.extend(FRAME_HEADER.pack(len(topic_raw), len(payload)))
        frames.extend(topic_raw)
        frames.extend(payload)
    capture.write_bytes(frames)
    metadata = [
        {
            "record_type": "message",
            "sequence": 1,
            "received_at_epoch_ns": 1_000_000_000,
            "received_monotonic_ns": 1_000_000_000,
        },
        {
            "record_type": "message",
            "sequence": 2,
            "received_at_epoch_ns": 11_000_000_000,
            "received_monotonic_ns": 11_000_000_000,
        },
    ]
    (tmp_path / "mqtt.metadata.jsonl").write_text(
        "".join(json.dumps(item) + "\n" for item in metadata),
        encoding="utf-8",
    )

    recording = FakeRecording()
    created: list[RerunBridge] = []

    def bridge_factory(**kwargs: object) -> RerunBridge:
        bridge = RerunBridge(
            recording_factory=lambda _: recording,  # type: ignore[arg-type]
            **kwargs,  # type: ignore[arg-type]
        )
        created.append(bridge)
        return bridge

    runtime = VisualizationRuntime(
        bridge_factory=bridge_factory,
    )
    runtime.start_replay(capture, speed=1.0)
    deadline = time.monotonic() + 5.0
    snapshot = runtime.snapshot()
    while time.monotonic() < deadline:
        snapshot = runtime.snapshot()
        if snapshot["rerun_grpc_url"] and snapshot["metrics"]["pcl_frames"] == 1:
            break
        time.sleep(0.05)

    assert snapshot["rerun_grpc_url"] == "rerun+http://127.0.0.1:9876/proxy"
    assert snapshot["metrics"]["pcl_frames"] == 1
    assert snapshot["metrics"]["mqtt_to_publish_ms"] is None
    runtime.stop(wait_seconds=5.0)
    assert runtime.snapshot()["phase"] == "idle"
    assert runtime.snapshot()["rerun_grpc_url"] == "rerun+http://127.0.0.1:9876/proxy"
    assert recording.disconnected is False

    runtime.start_replay(capture, speed=0.0)
    deadline = time.monotonic() + 5.0
    while time.monotonic() < deadline:
        if runtime.snapshot()["metrics"]["pcl_frames"] == 1:
            break
        time.sleep(0.05)
    runtime.stop(wait_seconds=5.0)

    assert len(created) == 1
    assert runtime.snapshot()["metrics"]["pcl_frames"] == 1
    runtime.close()
    assert runtime.snapshot()["rerun_grpc_url"] is None
    assert recording.disconnected is True


def test_close_during_blocked_factory_closes_the_late_bridge(tmp_path: Path) -> None:
    capture = tmp_path / "mqtt.raw.k1mqtt"
    topic = b"RealtimePointcloud"
    payload = struct.pack("<III", 16, 0, 0) + struct.pack(
        "<fffBBBB", 1.0, -2.0, 3.0, 10, 20, 30, 40
    )
    capture.write_bytes(RAW_MAGIC + FRAME_HEADER.pack(len(topic), len(payload)) + topic + payload)
    (tmp_path / "mqtt.metadata.jsonl").write_text(
        json.dumps(
            {
                "record_type": "message",
                "sequence": 1,
                "received_at_epoch_ns": 1_000_000_000,
                "received_monotonic_ns": 1_000_000_000,
            }
        )
        + "\n",
        encoding="utf-8",
    )

    factory_entered = threading.Event()
    release_factory = threading.Event()
    recording = FakeRecording()

    def blocked_factory(**kwargs: object) -> RerunBridge:
        factory_entered.set()
        assert release_factory.wait(timeout=5.0)
        return RerunBridge(
            recording_factory=lambda _: recording,  # type: ignore[arg-type]
            **kwargs,  # type: ignore[arg-type]
        )

    runtime = VisualizationRuntime(bridge_factory=blocked_factory)
    runtime.start_replay(capture, speed=0.0)
    assert factory_entered.wait(timeout=2.0)

    runtime.close(wait_seconds=0.01)
    release_factory.set()
    deadline = time.monotonic() + 2.0
    while time.monotonic() < deadline and not recording.disconnected:
        time.sleep(0.01)
    runtime.close(wait_seconds=1.0)

    assert recording.disconnected is True
    assert runtime.snapshot()["rerun_grpc_url"] is None
    with pytest.raises(RuntimeError, match="runtime завершён"):
        runtime.start_replay(capture, speed=0.0)


def test_stop_fails_closed_when_runtime_thread_misses_deadline(tmp_path: Path) -> None:
    capture = tmp_path / "mqtt.raw.k1mqtt"
    topic = b"RealtimePointcloud"
    payload = struct.pack("<III", 16, 0, 0) + struct.pack(
        "<fffBBBB", 1.0, -2.0, 3.0, 10, 20, 30, 40
    )
    capture.write_bytes(RAW_MAGIC + FRAME_HEADER.pack(len(topic), len(payload)) + topic + payload)
    (tmp_path / "mqtt.metadata.jsonl").write_text(
        json.dumps(
            {
                "record_type": "message",
                "sequence": 1,
                "received_at_epoch_ns": 1_000_000_000,
                "received_monotonic_ns": 1_000_000_000,
            }
        )
        + "\n",
        encoding="utf-8",
    )

    factory_entered = threading.Event()
    release_factory = threading.Event()

    def blocked_factory(**kwargs: object) -> RerunBridge:
        factory_entered.set()
        assert release_factory.wait(timeout=5.0)
        return RerunBridge(
            recording_factory=lambda _: FakeRecording(),  # type: ignore[arg-type]
            **kwargs,  # type: ignore[arg-type]
        )

    runtime = VisualizationRuntime(bridge_factory=blocked_factory)
    runtime.start_replay(capture, speed=0.0)
    assert factory_entered.wait(timeout=2.0)

    with pytest.raises(RuntimeError, match="не завершился"):
        runtime.stop(wait_seconds=0.01)
    assert runtime.snapshot()["phase"] == "stopping"

    release_factory.set()
    runtime.stop(wait_seconds=2.0)
    assert runtime.snapshot()["source_mode"] == "idle"
    runtime.close()