NODEDC_MISSION_CORE/tests/test_live_perception_shadow.py

from __future__ import annotations

import asyncio
import stat
import threading
import time
from pathlib import Path

import numpy as np

from k1link.compute.live_perception import (
    LivePerceptionIngress,
    encode_live_perception_result,
)
from k1link.device_plugins.xgrids_k1.live_perception_shadow import (
    build_live_perception_shadow_router,
    ensure_live_shadow_token,
)


def test_shadow_token_is_stable_and_private(tmp_path: Path) -> None:
    path, token = ensure_live_shadow_token(tmp_path)
    repeated_path, repeated_token = ensure_live_shadow_token(tmp_path)

    assert repeated_path == path
    assert repeated_token == token
    assert stat.S_IMODE(path.stat().st_mode) == 0o600
    assert stat.S_IMODE(path.parent.stat().st_mode) == 0o700


def test_shadow_router_exposes_only_the_exclusive_binary_stream() -> None:
    ingress = LivePerceptionIngress()
    router = build_live_perception_shadow_router(
        ingress,
        "test-plugin",
        bearer_token="x" * 43,
    )
    assert len(router.routes) == 1
    assert router.routes[0].path == (
        "/api/v1/device-plugins/test-plugin/live-perception-shadow"
    )


def test_shadow_router_accepts_only_validated_diagnostic_results_back() -> None:
    ingress = LivePerceptionIngress()
    received: list[bytes] = []
    published = threading.Event()

    def receive_result(payload: bytes) -> bool:
        received.append(payload)
        published.set()
        return True

    router = build_live_perception_shadow_router(
        ingress,
        "test-plugin",
        bearer_token="x" * 43,
        result_receiver=receive_result,
    )
    ingress.begin_session("session-1")
    encoded = encode_live_perception_result(
        frame_index=0,
        source_frame_index=0,
        session_seconds=0.0,
        captured_at_epoch_ns=1,
        image_jpeg=bytes.fromhex("ffd878ffd9"),
        segmentation_mask=np.zeros((600, 800), dtype=np.uint8),
        objects=[],
        delivery={"health": "healthy"},
    )

    class FakeWebSocket:
        def __init__(self) -> None:
            self.headers = {"authorization": f"Bearer {'x' * 43}"}
            self.messages = [
                {"type": "websocket.receive", "bytes": encoded},
                {"type": "websocket.disconnect"},
            ]
            self.sent: list[bytes] = []
            self.accepted = False

        async def accept(self) -> None:
            self.accepted = True

        async def receive(self) -> dict[str, object]:
            await asyncio.sleep(0)
            return self.messages.pop(0)

        async def send_bytes(self, payload: bytes) -> None:
            self.sent.append(payload)

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

    websocket = FakeWebSocket()
    asyncio.run(router.routes[0].endpoint(websocket))  # type: ignore[attr-defined]

    assert websocket.accepted is True
    assert published.is_set()
    assert received == [encoded]


def test_shadow_router_does_not_discard_ingress_while_receiving_results() -> None:
    class SlowIngress(LivePerceptionIngress):
        def take_next(
            self,
            consumer_id: str,
            *,
            timeout: float | None = None,
        ):
            time.sleep(0.01)
            return super().take_next(consumer_id, timeout=timeout)

    ingress = SlowIngress()
    received: list[bytes] = []
    router = build_live_perception_shadow_router(
        ingress,
        "test-plugin",
        bearer_token="x" * 43,
        result_receiver=lambda payload: not received.append(payload),
    )
    ingress.begin_session("session-1")
    for modality, sequence in (
        ("camera-init", 0),
        ("camera-frame", 1),
        ("lidar", 1),
        ("pose", 1),
    ):
        assert ingress.publish(
            modality=modality,
            source_id="sensor.camera.right",
            source_sequence=sequence,
            captured_at_epoch_ns=sequence,
            received_monotonic_ns=sequence,
            payload=modality.encode(),
        )
    ingress.end_session("session-1")
    expected_events = sum(
        int(queue["depth"]) for queue in ingress.snapshot()["queues"].values()
    )

    class DuplexFakeWebSocket:
        def __init__(self) -> None:
            self.headers = {"authorization": f"Bearer {'x' * 43}"}
            self.results_remaining = 8
            self.sent: list[bytes] = []

        async def accept(self) -> None:
            return

        async def receive(self) -> dict[str, object]:
            if self.results_remaining:
                self.results_remaining -= 1
                await asyncio.sleep(0)
                return {"type": "websocket.receive", "bytes": b"result"}
            deadline = asyncio.get_running_loop().time() + 1
            while len(self.sent) < expected_events:
                if asyncio.get_running_loop().time() >= deadline:
                    break
                await asyncio.sleep(0.005)
            return {"type": "websocket.disconnect"}

        async def send_bytes(self, payload: bytes) -> None:
            self.sent.append(payload)

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

    websocket = DuplexFakeWebSocket()
    asyncio.run(router.routes[0].endpoint(websocket))  # type: ignore[attr-defined]

    assert len(received) == 8
    assert len(websocket.sent) == expected_events
    snapshot = ingress.snapshot()
    assert sum(
        int(queue["consumed"]) for queue in snapshot["queues"].values()
    ) == expected_events