NODEDC_MISSION_CORE/tests/test_session_lab_cache.py

from __future__ import annotations

import hashlib
import json
from pathlib import Path

from k1link.sessions.lab_cache import publish_lab_replay_cache
from k1link.sessions.recording import (
    CACHE_SCHEMA,
    RECORDING_CACHE_FILENAME,
    RECORDING_CACHE_SIDECAR_FILENAME,
)


def canonical_json(value: object) -> bytes:
    return json.dumps(
        value,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
        allow_nan=False,
    ).encode()


def test_publish_lab_replay_cache_hardlinks_rrd_and_rebinds_sidecars(
    tmp_path: Path,
) -> None:
    data = tmp_path / "mission-core"
    recordings = data / "recordings"
    source = recordings / "source-session"
    source.mkdir(parents=True)
    source_rrd = source / RECORDING_CACHE_FILENAME
    source_rrd.write_bytes(b"RRF2immutable")
    source_sidecar = {
        "schema_version": CACHE_SCHEMA,
        "session_id": source.name,
        "recording_byte_length": source_rrd.stat().st_size,
        "recording_mtime_ns": source_rrd.stat().st_mtime_ns,
        "recording_sha256": hashlib.sha256(source_rrd.read_bytes()).hexdigest(),
        "timeline": "session_time",
        "timeline_start_ns": 0,
        "timeline_end_ns": 1,
    }
    (source / RECORDING_CACHE_SIDECAR_FILENAME).write_text(
        json.dumps(source_sidecar),
        encoding="utf-8",
    )
    media_root = data / "recorded-media-preparations"
    media_root.mkdir()
    media_body = {
        "schema_version": "missioncore.recorded-media-preparation/v1",
        "session_id": source.name,
        "artifact_id": "recorded-video-1",
        "generation": "accepted",
    }
    (media_root / "source.json").write_text(
        json.dumps(
            {
                **media_body,
                "checksum_sha256": hashlib.sha256(canonical_json(media_body)).hexdigest(),
            }
        ),
        encoding="utf-8",
    )

    publish_lab_replay_cache(
        data,
        source_session_id=source.name,
        lab_session_id="lab-e21-result",
    )

    lab = recordings / "lab-e21-result"
    lab_rrd = lab / RECORDING_CACHE_FILENAME
    assert lab_rrd.read_bytes() == source_rrd.read_bytes()
    assert lab_rrd.stat().st_ino == source_rrd.stat().st_ino
    assert lab_rrd.stat().st_nlink == 2
    cloned_cache = json.loads(
        (lab / RECORDING_CACHE_SIDECAR_FILENAME).read_text(encoding="utf-8")
    )
    assert cloned_cache["session_id"] == "lab-e21-result"
    cloned_media = [
        json.loads(path.read_text(encoding="utf-8"))
        for path in media_root.iterdir()
        if path.name != "source.json"
    ]
    assert len(cloned_media) == 1
    assert cloned_media[0]["session_id"] == "lab-e21-result"
    cloned_checksum = cloned_media[0].pop("checksum_sha256")
    assert cloned_checksum == hashlib.sha256(canonical_json(cloned_media[0])).hexdigest()

    lab_rrd.unlink()
    assert source_rrd.read_bytes() == b"RRF2immutable"


def test_publish_lab_replay_cache_materializes_exact_bounded_window(
    tmp_path: Path,
) -> None:
    import numpy as np
    import rerun as rr

    data = tmp_path / "mission-core"
    recordings = data / "recordings"
    source = recordings / "source-session"
    source.mkdir(parents=True)
    source_rrd = source / RECORDING_CACHE_FILENAME
    recording = rr.RecordingStream(
        "nodedc_mission_core_recorded",
        recording_id=source.name,
    )
    stream = rr.binary_stream(recording)
    try:
        for index, timestamp in enumerate((0, 2_000_000_000, 4_000_000_000, 8_000_000_000)):
            recording.set_time(
                "session_time",
                duration=np.timedelta64(timestamp, "ns"),
            )
            recording.log(
                "/world/points",
                rr.Points3D([[float(index), 0.0, 0.0]]),
            )
        payload = stream.read(flush=True, flush_timeout_sec=5.0)
    finally:
        recording.disconnect()
    assert payload is not None
    source_rrd.write_bytes(payload)
    source_sidecar = {
        "schema_version": CACHE_SCHEMA,
        "session_id": source.name,
        "recording_byte_length": source_rrd.stat().st_size,
        "recording_mtime_ns": source_rrd.stat().st_mtime_ns,
        "recording_sha256": hashlib.sha256(source_rrd.read_bytes()).hexdigest(),
        "timeline": "session_time",
        "timeline_start_ns": 0,
        "timeline_end_ns": 8_000_000_000,
    }
    (source / RECORDING_CACHE_SIDECAR_FILENAME).write_text(
        json.dumps(source_sidecar),
        encoding="utf-8",
    )

    publish_lab_replay_cache(
        data,
        source_session_id=source.name,
        lab_session_id="lab-e21-bounded",
        timeline_start_ns=2_000_000_000,
        timeline_end_ns=4_000_000_000,
    )

    lab = recordings / "lab-e21-bounded"
    lab_rrd = lab / RECORDING_CACHE_FILENAME
    assert lab_rrd.stat().st_ino != source_rrd.stat().st_ino
    bounded = json.loads(
        (lab / RECORDING_CACHE_SIDECAR_FILENAME).read_text(encoding="utf-8")
    )
    assert bounded["timeline_start_ns"] == 2_000_000_000
    assert bounded["timeline_end_ns"] == 4_000_000_000
    assert bounded["recording_byte_length"] == lab_rrd.stat().st_size
    assert bounded["recording_sha256"] == hashlib.sha256(lab_rrd.read_bytes()).hexdigest()
    assert "lab_window" not in bounded

    original_stat = lab_rrd.stat()
    publish_lab_replay_cache(
        data,
        source_session_id=source.name,
        lab_session_id="lab-e21-bounded",
        timeline_start_ns=2_000_000_000,
        timeline_end_ns=4_000_000_000,
    )
    assert lab_rrd.stat().st_ino == original_stat.st_ino
    assert lab_rrd.stat().st_mtime_ns == original_stat.st_mtime_ns