orisys-sdk/tests/test_pipeline.py

"""StreamPipeline: staging, snapshot immutability, error propagation, shutdown."""

from __future__ import annotations

import threading
import time

import numpy as np
import pytest

from orisys import FrameSnapshot, StreamPipeline


class FakeSensor:
    """Minimal Sensor stand-in with the same mutable-buffer behaviour.

    It deliberately reuses one image buffer (like ``Sensor._img_gray_buf``) so the
    pipeline's private-copy contract is actually exercised.
    """

    class info:
        VRAW = "vraw"
        VNORMAL = "vnormal"
        VSHEAR = "vshear"
        DEPTH = "depth"
        CENTROID = "centroid"
        CONTOUR = "contour"

    def __init__(self, *, fail_at: int | None = None, flow_delay: float = 0.0,
                 hhd_delay: float = 0.0):
        self.img = None
        self.flow = None
        self.vnormal = None
        self.vshear = None
        self.depth_map = None
        self.centroid = (0, 0)
        self.contour = np.empty(0)
        self.ftor = 0.0
        self._force_raw = (0.0, 0.0, 0.0)
        self.fnormal = self.fshearx = self.fsheary = 0.0
        self.force_unit = "raw"
        self._img_buf = np.zeros((8, 8), dtype=np.uint8)
        self._flow_buf = np.zeros((8, 8, 2), dtype=np.float32)
        self._flow_delay = flow_delay
        self._hhd_delay = hhd_delay
        self._fail_at = fail_at
        self._flow_calls = 0
        self._seen_images: list[np.ndarray] = []
        self._seen_flow_values: list[float] = []
        self.contact_calls = 0
        #: Mirrors ``Sensor.reset()``, which rebases ``initial_img`` on the current frame. Records
        #: the frame sequence at each call so a test can prove the reset did not race the flow
        #: stage's read of that buffer.
        self.img_size_raw = None
        self.reset_calls = 0
        self.reset_at_frame: list[int] = []

    def make_frame(self, seq: int) -> np.ndarray:
        # writes INTO the shared buffer, exactly like cv2.resize(dst=...)
        self._img_buf.fill(seq % 251)
        self.img = self._img_buf
        return self._img_buf

    def compute_flow(self, img=None, **kwargs) -> None:
        image = self.img if img is None else img
        self._flow_calls += 1
        if self._fail_at is not None and self._flow_calls == self._fail_at:
            raise RuntimeError("flow exploded")
        self._seen_images.append(np.array(image, copy=True))
        time.sleep(self._flow_delay)
        # rewrite the shared buffer, like the real DIS/EMA path does
        self._flow_buf.fill(float(self._flow_calls))
        self.flow = self._flow_buf

    def compute_hhd(self, flow=None, **kwargs) -> None:
        # Mirrors Sensor.compute_hhd, which derives the forces from the field it
        # is given - so this only yields frame N's forces if the pipeline passes
        # frame N's flow explicitly.
        time.sleep(self._hhd_delay)
        field = self.flow if flow is None else flow
        value = (float(field[0, 0, 0])
                 if isinstance(field, np.ndarray) and field.size else 0.0)
        self._seen_flow_values.append(value)
        self.fnormal = value * 10.0
        self.fshearx = value
        self.fsheary = -value
        self._force_raw = (self.fnormal, self.fshearx, self.fsheary)
        self.ftor = value
        self.vnormal = self._flow_buf
        self.vshear = self._flow_buf

    def compute_contact(self) -> None:
        self.contact_calls += 1
        self.centroid = (1, 2)
        self.contour = np.zeros((1, 1, 2), dtype=np.int32)
        self.depth_map = np.zeros((8, 8), dtype=np.float32)

    def reset(self) -> None:
        self.reset_calls += 1
        self.reset_at_frame.append(self._flow_calls)

    def read_info(self, *args):
        source = {
            self.info.VRAW: lambda: self.flow.copy(),
            self.info.VNORMAL: lambda: self.vnormal.copy(),
            self.info.VSHEAR: lambda: self.vshear.copy(),
            self.info.DEPTH: lambda: self.depth_map.copy(),
            self.info.CENTROID: lambda: np.array(self.centroid),
            self.info.CONTOUR: lambda: self.contour.copy(),
        }
        return tuple(source[a]() for a in args)


def make_producer(sensor: FakeSensor, count: int):
    state = {"i": 0}

    def produce():
        if state["i"] >= count:
            return None
        frame = sensor.make_frame(state["i"])
        state["i"] += 1
        return frame

    return produce


def run_all(sensor: FakeSensor, count: int, **kwargs):
    """Drive the pipeline with an exhaustive consumer: every frame, in order."""
    frames: list = []
    kwargs.setdefault("depth", 2)
    pipe = StreamPipeline(sensor, make_producer(sensor, count),
                          on_frame=frames.append, **kwargs)
    pipe.start()
    deadline = time.perf_counter() + 10.0
    while len(frames) < count and time.perf_counter() < deadline:
        if pipe.error is not None:
            break
        time.sleep(0.001)
    pipe.stop()
    return pipe, frames


def test_on_frame_receives_every_frame_in_order():
    sensor = FakeSensor()
    pipe, frames = run_all(sensor, 12)
    assert len(frames) == 12
    assert [s.seq for s in frames] == list(range(12))
    assert pipe.stats.produced == 12
    assert pipe.stats.published == 12


def test_forces_come_from_the_matching_frame():
    """Each snapshot must carry the forces computed from its own frame's flow."""
    sensor = FakeSensor()
    _, frames = run_all(sensor, 10)
    # frame i's flow value is (i + 1); the fake derives the force from self.flow,
    # so a mismatched pairing would show up as a repeated or skipped value.
    assert sensor._seen_flow_values == [float(i + 1) for i in range(10)]
    assert [s.normal_force for s in frames] == [10.0 * (i + 1) for i in range(10)]
    assert [s.shear_x for s in frames] == [float(i + 1) for i in range(10)]


def test_stages_do_not_read_a_rewritten_image_buffer():
    """The producer reuses one buffer; no stage may observe a later frame."""
    sensor = FakeSensor(flow_delay=0.004)
    run_all(sensor, 8)
    seen = sensor._seen_images
    assert len(seen) == 8
    for i, arr in enumerate(seen):
        assert np.all(arr == i % 251), f"frame {i} was overwritten before flow ran"


def test_snapshot_arrays_are_private_copies():
    sensor = FakeSensor()
    _, frames = run_all(sensor, 4, include=("image", "flow", "vnormal"))
    first = frames[0]
    assert isinstance(first, FrameSnapshot)
    assert first.image is not None and first.flow is not None
    saved = first.flow.copy()
    # the SDK keeps rewriting its buffers; the snapshot must not change
    sensor._flow_buf.fill(999.0)
    sensor._img_buf.fill(7)
    assert np.array_equal(first.flow, saved)
    assert not np.any(first.flow == 999.0)


def test_include_filters_which_arrays_are_copied():
    sensor = FakeSensor()
    _, frames = run_all(sensor, 3, include=("vnormal",))
    assert frames[0].vnormal is not None
    assert frames[0].image is None and frames[0].flow is None
    assert frames[0].depth is None and frames[0].contour is None


def test_rejects_unknown_include_field():
    sensor = FakeSensor()
    with pytest.raises(ValueError):
        StreamPipeline(sensor, make_producer(sensor, 1), include=("floww",))


def test_contact_stage_runs_when_requested():
    sensor = FakeSensor()
    _, frames = run_all(sensor, 5, contact=True,
                        include=("depth", "centroid", "contour"))
    assert sensor.contact_calls == 5
    assert frames[-1].depth is not None
    assert frames[-1].centroid is not None
    assert frames[-1].contour is not None


def test_stage_exception_surfaces_on_stop():
    sensor = FakeSensor(fail_at=3)
    pipe = StreamPipeline(sensor, make_producer(sensor, 20), depth=2)
    pipe.start()
    time.sleep(0.4)
    with pytest.raises(RuntimeError, match="flow exploded"):
        pipe.stop()


def test_stage_exception_surfaces_on_wait():
    sensor = FakeSensor(fail_at=2)
    pipe = StreamPipeline(sensor, make_producer(sensor, 20), depth=2)
    pipe.start()
    with pytest.raises(RuntimeError):
        for _ in range(50):
            pipe.wait(timeout=0.1)
    pipe.stop()


def test_stop_is_idempotent():
    sensor = FakeSensor()
    pipe, _ = run_all(sensor, 4)
    pipe.stop()                     # second stop must not raise
    assert not pipe.running
    assert pipe.latest() is not None


def test_stop_while_running_does_not_hang():
    """Stopping mid-stream must unblock the consumer threads.

    The producer stops delivering sentinels once the done flag is set, so a
    bare blocking Queue.get() in a consumer would wait forever.
    """
    sensor = FakeSensor(flow_delay=0.002)
    pipe = StreamPipeline(sensor, make_producer(sensor, 100000), depth=2)
    pipe.start()
    assert pipe.wait(timeout=2.0) is not None
    t0 = time.perf_counter()
    pipe.stop(timeout=5.0)
    elapsed = time.perf_counter() - t0
    assert elapsed < 2.0, f"stop() took {elapsed:.2f}s; a stage was still blocked"
    assert not pipe.running


def test_context_manager_stops_threads():
    sensor = FakeSensor()
    frames: list = []
    with StreamPipeline(sensor, make_producer(sensor, 6),
                        on_frame=frames.append) as pipe:
        deadline = time.perf_counter() + 5.0
        while len(frames) < 6 and time.perf_counter() < deadline:
            time.sleep(0.001)
    assert len(frames) == 6
    assert not pipe.running


def test_wait_exposes_a_latest_wins_view():
    sensor = FakeSensor()
    pipe = StreamPipeline(sensor, make_producer(sensor, 5))
    pipe.start()
    last = None
    while True:
        snap = pipe.wait(timeout=2.0)
        if snap is None:
            break
        if last is not None:
            assert snap.seq > last.seq
        last = snap
    # the pipeline has drained; the final snapshot is the last frame
    assert last is not None and last.normal_force == 50.0
    assert pipe.latest() is last
    # nothing new is coming, so a further poll returns immediately
    assert pipe.wait(timeout=0.2) is None
    pipe.stop()


def test_wait_timeout_returns_none_without_new_frames():
    sensor = FakeSensor()
    pipe = StreamPipeline(sensor, make_producer(sensor, 1))
    pipe.start()
    assert pipe.wait(timeout=1.0) is not None
    t0 = time.perf_counter()
    assert pipe.wait(timeout=0.1) is None
    assert time.perf_counter() - t0 < 1.0
    pipe.stop()


def test_stats_track_stage_times_and_age():
    sensor = FakeSensor(flow_delay=0.002, hhd_delay=0.001)
    pipe, _ = run_all(sensor, 15)
    d = pipe.stats.as_dict()
    assert d["published"] == 15
    assert d["flow_ms"] >= 1.5           # the injected delay shows up
    assert d["hhd_ms"] >= 0.8
    assert d["age_ms"] > 0.0
    assert d["fps"] > 0.0
    assert d["produce_ms"] >= 0.0


def test_depth_does_not_change_results():
    for depth in (1, 2, 4):
        sensor = FakeSensor()
        _, frames = run_all(sensor, 10, depth=depth)
        assert [s.normal_force for s in frames] == [10.0 * (i + 1) for i in range(10)]


def test_for_finger_runtime_maps_eof_to_none():
    """``for_finger_runtime`` drives ``grab_and_process`` and maps EOF to None."""
    sensor = FakeSensor()
    state = {"i": 0}

    class FakeRuntime:
        tactile = type("T", (), {"raw_sensor": sensor})()

        def grab_and_process(self, color=False):
            if state["i"] >= 3:
                return False, None
            frame = sensor.make_frame(state["i"])
            state["i"] += 1
            return True, frame

    frames: list = []
    pipe = StreamPipeline.for_finger_runtime(FakeRuntime(), on_frame=frames.append)
    pipe.start()
    deadline = time.perf_counter() + 5.0
    while len(frames) < 3 and time.perf_counter() < deadline:
        time.sleep(0.001)
    pipe.stop()
    assert len(frames) == 3
    assert pipe.stats.produced == 3


def test_reset_is_applied_on_the_flow_stage_not_the_caller():
    """``reset()`` must not run on the caller's thread.

    ``Sensor.reset`` rebases ``initial_img`` on the current ``Sensor.img``, and the flow stage is
    reading that buffer concurrently, so a caller-thread reset would capture a half-written
    reference frame. The pipeline queues it and applies it before the next frame instead.
    """
    sensor = FakeSensor()
    ready = threading.Event()
    release = threading.Event()

    def gated_produce():
        # Hold the producer until the caller has asked for a reset, so the reset is provably
        # queued before any frame reaches the flow stage.
        ready.set()
        release.wait(timeout=5.0)
        return sensor.make_frame(0)

    frames: list = []
    pipe = StreamPipeline(sensor, gated_produce, on_frame=frames.append)
    pipe.start()
    assert ready.wait(timeout=5.0)

    pipe.reset()
    pipe.reset()          # idempotent: a second request before the next frame is not doubled
    assert sensor.reset_calls == 0, "reset must not run on the caller's thread"

    release.set()
    deadline = time.perf_counter() + 5.0
    while not frames and time.perf_counter() < deadline:
        time.sleep(0.001)
    pipe.stop()

    assert sensor.reset_calls == 1, f"expected one queued reset, got {sensor.reset_calls}"
    assert sensor.reset_at_frame == [0], (
        "the reset must land before the first compute_flow, not after it"
    )


def test_snapshot_carries_the_raw_capture_size():
    """``image_size_raw`` lets a caller compare the processed frame with the capture resolution."""
    sensor = FakeSensor()
    sensor.img_size_raw = (1280, 720)
    frames: list = []
    pipe = StreamPipeline(sensor, make_producer(sensor, 1), on_frame=frames.append)
    pipe.start()
    deadline = time.perf_counter() + 5.0
    while not frames and time.perf_counter() < deadline:
        time.sleep(0.001)
    pipe.stop()
    assert frames and frames[0].image_size_raw == (1280, 720)


def test_snapshot_raw_capture_size_is_none_when_unreported():
    """A producer that never sets the size must not fabricate one."""
    sensor = FakeSensor()
    sensor.img_size_raw = None
    frames: list = []
    pipe = StreamPipeline(sensor, make_producer(sensor, 1), on_frame=frames.append)
    pipe.start()
    deadline = time.perf_counter() + 5.0
    while not frames and time.perf_counter() < deadline:
        time.sleep(0.001)
    pipe.stop()
    assert frames and frames[0].image_size_raw is None