orisys-app/shared/tests/test_depth_replay.py

"""Offline depth replay: where to start, how far to step, which chain to run.

P1-3 and P1-4: the pane kept showing the previous frame's depth after the cursor
moved, the recomputer re-read ``target`` frames on every call (so requesting
frames 2 and 3 landed on 5), it skipped the live processing chain, and it
reduced the recorded config path to a name that could load a different file.
"""

from __future__ import annotations

import shutil
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch

import numpy as np
import pytest

from orisys_app_shared.record.depth import DepthRecomputer

SCRATCH = Path(__file__).resolve().parent / "_scratch"


def _root() -> Path:
    shutil.rmtree(SCRATCH, ignore_errors=True)
    SCRATCH.mkdir(parents=True, exist_ok=True)
    return SCRATCH


class _CountingSensor:
    """A sensor whose reported depth is the number of frames it has read."""

    opened: list[dict] = []

    def __init__(self, *args, **kwargs):
        _CountingSensor.opened.append({"args": args, "kwargs": kwargs})
        self.n = 0
        self.chain_calls = 0
        self.info = SimpleNamespace(DEPTH=1)

    def get_img(self):
        self.n += 1
        return np.ones((2, 2))

    def compute_deformation(self, **_kwargs):
        self.chain_calls += 1

    def compute_contact(self):
        self.chain_calls += 1

    def read_info(self, _channel):
        return np.array([self.n])

    def disconnect(self):
        return None


def _reader(session: dict, *, segments=("raw_0001.mp4",), frames=()):
    directory = _root()
    for name in segments:
        (directory / name).write_bytes(b"")
    return SimpleNamespace(
        session=lambda: session,
        segment_path=lambda segment: directory / f"raw_{segment:04d}.mp4",
        segments=lambda: [directory / name for name in segments],
        info=lambda: SimpleNamespace(frames=list(frames)),
        frame_of=lambda frame_id: next(
            (frame for frame in frames if frame.frame_id == frame_id), None
        ),
        frame_at=lambda t: None,
        last_reset_frame=lambda frame_id: None,
        segment_start_frame=lambda segment: next(
            (frame for frame in frames if frame.segment == segment), None
        ),
        preroll_frames=lambda: 0,
        preroll_segment=lambda: None,
        directory=directory,
    )


def _session(config_ref: str) -> dict:
    return {"device": {"config_ref": config_ref, "sensor_type": "box"}}


def _config_file() -> str:
    """A config that exists, so ``preflight`` is not what a test measures."""
    path = _root() / "sensor.json"
    path.write_text("{}", encoding="utf-8")
    return str(path)


def _reader_with_config(frames):
    """A reader whose recorded config really exists on disk.

    ``_reader`` recreates the scratch directory, so the config file is written
    *after* the reader is built.
    """
    directory = _root()
    config = directory / "sensor.json"
    reader = _reader(_session(str(config)), frames=frames)
    config.write_text("{}", encoding="utf-8")
    return reader


def test_stepping_continues_from_the_current_position():
    """The old code re-read ``target`` frames every call: 2 then 3 landed on 5."""
    _CountingSensor.opened.clear()
    computer = DepthRecomputer(_reader(_session("/cfg/sensor.json")))
    with patch("orisys.Sensor", _CountingSensor):
        first = computer._depth_from_segment(Path("raw_0001.mp4"), 2)
        second = computer._depth_from_segment(Path("raw_0001.mp4"), 3)
    assert first.tolist() == [2]
    assert second.tolist() == [3]
    # One sensor for the whole replay: reopening would reset the algorithm state.
    assert len(_CountingSensor.opened) == 1
    computer.close()


def test_stepping_backwards_reopens_the_segment():
    _CountingSensor.opened.clear()
    computer = DepthRecomputer(_reader(_session("/cfg/sensor.json")))
    with patch("orisys.Sensor", _CountingSensor):
        assert computer._depth_from_segment(Path("raw_0001.mp4"), 5).tolist() == [5]
        assert computer._depth_from_segment(Path("raw_0001.mp4"), 2).tolist() == [2]
    assert len(_CountingSensor.opened) == 2, "seeking backwards needs a fresh source"
    computer.close()


def test_the_live_processing_chain_runs_for_every_replayed_frame():
    """Without it the replay is a different algorithm than the recording."""
    _CountingSensor.opened.clear()
    computer = DepthRecomputer(_reader(_session("/cfg/sensor.json")))
    with patch("orisys.Sensor", _CountingSensor):
        computer._depth_from_segment(Path("raw_0001.mp4"), 3)
        sensor = computer._sensor._sensor
    assert sensor.chain_calls == 6, "one deformation + one contact per frame"
    computer.close()


def test_the_exact_config_path_is_opened(tmp_path=None):
    """A reduced name could load a different file with the same name."""
    _CountingSensor.opened.clear()
    config = Path("D:/custom/box/sensor.json")
    computer = DepthRecomputer(_reader(_session(str(config))))
    with patch("orisys.Sensor", _CountingSensor):
        computer._depth_from_segment(Path("raw_0001.mp4"), 1)
    assert _CountingSensor.opened[0]["kwargs"]["config_name"] == str(config)
    computer.close()


def test_the_plan_starts_at_the_last_reset_in_the_same_segment():
    frames = [
        SimpleNamespace(
            frame_id=index,
            t_rel_s=index / 30.0,
            segment=1,
            segment_frame=index + 1,
            in_video=True,
            dropped=False,
        )
        for index in range(10)
    ]
    reader = _reader(_session("/cfg/sensor.json"), frames=frames)
    reader.last_reset_frame = lambda frame_id: 6
    computer = DepthRecomputer(reader)
    plan = computer.plan_for(frames[8])
    assert plan.stop_at == 9
    assert plan.preroll is None
    assert plan.warnings == ()
    computer.close()


def test_the_plan_uses_the_pre_roll_for_the_start_of_a_recording():
    frames = [
        SimpleNamespace(
            frame_id=index,
            t_rel_s=index / 30.0,
            segment=1,
            segment_frame=index + 1,
            in_video=True,
            dropped=False,
        )
        for index in range(5)
    ]

    def reader_factory():
        reader = _reader(_session("/cfg/sensor.json"), frames=frames)
        reader.preroll_frames = lambda: 30
        reader.preroll_segment = lambda: reader.directory / "preroll_0001.mp4"
        return reader

    reader = reader_factory()
    (reader.directory / "preroll_0001.mp4").write_bytes(b"")
    computer = DepthRecomputer(reader)
    plan = computer.plan_for(frames[3])
    assert plan.preroll is not None
    assert plan.preroll_frames == 30
    assert plan.warnings == ()
    computer.close()


def test_a_segment_without_history_is_flagged_approximate():
    """Starting mid-recording with no reset and no pre-roll cannot be exact."""
    frames = [
        SimpleNamespace(
            frame_id=index,
            t_rel_s=index / 30.0,
            segment=2,
            segment_frame=index + 1,
            in_video=True,
            dropped=False,
        )
        for index in range(5)
    ]
    reader = _reader(_session("/cfg/sensor.json"), frames=frames)
    computer = DepthRecomputer(reader)
    plan = computer.plan_for(frames[2])
    assert "state_approximate" in plan.warnings
    computer.close()


def test_a_dropped_frame_is_reported_not_substituted():
    """The index marks the hole; replay must say so instead of showing a neighbour."""
    frames = [
        SimpleNamespace(
            frame_id=0,
            t_rel_s=0.0,
            segment=1,
            segment_frame=1,
            in_video=True,
            dropped=False,
        ),
        SimpleNamespace(
            frame_id=1,
            t_rel_s=1 / 30.0,
            segment=1,
            segment_frame=0,
            in_video=False,
            dropped=True,
        ),
    ]
    reader = _reader_with_config(frames)
    reader.frame_at = lambda _t: frames[1]
    computer = DepthRecomputer(reader)
    result = computer.depth_at(1 / 30.0)
    assert not result.ok
    assert result.error == "frame_not_recorded"
    computer.close()