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, ProducedFrame, 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_producer_metadata_stays_with_its_image_through_the_queues():
    sensor = FakeSensor(flow_delay=0.002)
    frames = []
    produced = 0

    def produce():
        nonlocal produced
        if produced == 12:
            return None
        value = produced
        produced += 1
        return ProducedFrame(sensor.make_frame(value), {"source_seq": value})

    pipe = StreamPipeline(sensor, produce, on_frame=frames.append)
    pipe.start()
    deadline = time.perf_counter() + 10.0
    while len(frames) < 12 and time.perf_counter() < deadline:
        time.sleep(0.001)
    pipe.stop()
    assert len(frames) == 12
    for snap in frames:
        assert snap.source_metadata["source_seq"] == int(snap.image[0, 0])


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


class SplitResetFakeSensor(FakeSensor):
    """A fake with the split reset seam and event-gated flow/HHD stages.

    The two halves of a reset belong to different pipeline threads, and the
    version a frame is stamped with depends on which half has already run. The
    events pin both stages to a known frame so the regression orders the race
    instead of hoping for it: the flow stage parks at the end of its second
    frame and the HHD stage parks inside its first decomposition.
    """

    FLOW_GATE_CALL = 2

    def __init__(self) -> None:
        super().__init__()
        self.flow_gate_reached = threading.Event()
        self.release_flow_gate = threading.Event()
        self.hhd_gate_reached = threading.Event()
        self.release_hhd_gate = threading.Event()
        self.hhd_calls = 0
        self.image_reset_threads: list[str] = []
        self.force_reset_threads: list[str] = []
        self.force_reset_at_hhd_calls: list[int] = []

    def compute_flow(self, img=None, **kwargs) -> None:
        super().compute_flow(img, **kwargs)
        # Park after frame 2's flow. The next frame that reaches the flow stage
        # is then provably the first one that can observe a reset queued now.
        if self._flow_calls == self.FLOW_GATE_CALL:
            self.flow_gate_reached.set()
            self.release_flow_gate.wait(timeout=5.0)

    def compute_hhd(self, flow=None, **kwargs) -> None:
        self.hhd_calls += 1
        if self.hhd_calls == 1:
            self.hhd_gate_reached.set()
            self.release_hhd_gate.wait(timeout=5.0)
        super().compute_hhd(flow, **kwargs)

    def reset_image_state(self) -> None:
        self.image_reset_threads.append(threading.current_thread().name)

    def reset_force_state(self) -> None:
        self.force_reset_threads.append(threading.current_thread().name)
        self.force_reset_at_hhd_calls.append(self.hhd_calls)

    def reset(self) -> None:
        # Mirrors Sensor.reset(): both halves, immediately, on the caller's thread.
        self.reset_calls += 1
        self.reset_at_frame.append(self._flow_calls)
        self.reset_image_state()
        self.reset_force_state()


class LegacyResetFakeSensor(FakeSensor):
    """A stand-in with only the historical ``reset()``, like a third-party sensor."""

    def __init__(self) -> None:
        super().__init__()
        self.reset_threads: list[str] = []

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


def test_reset_request_target_is_shared_until_flow_claims_it():
    sensor = FakeSensor()
    frames = []
    pipe = StreamPipeline(sensor, make_producer(sensor, 2), on_frame=frames.append)
    assert pipe.request_reset() == 1
    assert pipe.reset() is None  # The published entry keeps its return contract.
    assert pipe.request_reset() == 1
    assert pipe.tracker_reset_seq == 0
    pipe.start()
    try:
        assert _wait_until(lambda: len(frames) == 2)
        assert [snap.tracker_reset_seq for snap in frames] == [1, 1]
        assert sensor.reset_calls == 1
        assert pipe.request_reset() == 2
    finally:
        pipe.stop()


def test_reset_requested_while_previous_reset_is_running_gets_a_new_target():
    entered = threading.Event()
    release = threading.Event()

    class Sensor(LegacyResetFakeSensor):
        def reset(self):
            super().reset()
            if self.reset_calls == 1:
                entered.set()
                assert release.wait(timeout=5.0)

    sensor = Sensor()
    frames = []
    pipe = StreamPipeline(sensor, make_producer(sensor, 2), on_frame=frames.append)
    assert pipe.request_reset() == 1
    pipe.start()
    try:
        assert entered.wait(timeout=5.0)
        # Flow has claimed request 1, but has not finished applying it. The live
        # counter is still 0; request 2 must not be merged into request 1.
        assert pipe.tracker_reset_seq == 0
        assert pipe.request_reset() == 2
        assert pipe.request_reset() == 2
        release.set()
        assert _wait_until(lambda: len(frames) == 2)
        assert [snap.tracker_reset_seq for snap in frames] == [1, 2]
        assert sensor.reset_calls == 2
    finally:
        release.set()
        pipe.stop()


def _wait_until(predicate, timeout: float = 5.0) -> bool:
    """Poll a predicate; the events, not this poll, order the race."""
    deadline = time.perf_counter() + timeout
    while time.perf_counter() < deadline:
        if predicate():
            return True
        time.sleep(0.001)
    return predicate()


def _run_gated_reset(sensor: SplitResetFakeSensor, frames: list, count: int = 3) -> None:
    """Drive two pre-reset frames and one post-reset frame through the real loops."""
    pipe = StreamPipeline(sensor, make_producer(sensor, count),
                          on_frame=frames.append, depth=2)
    pipe.start()
    try:
        assert sensor.flow_gate_reached.wait(timeout=5.0), "flow stage never reached its gate"
        assert sensor.hhd_gate_reached.wait(timeout=5.0), "HHD stage never reached its gate"
        # Both stages are provably parked: no frame has been published, and the
        # frame whose flow would observe this reset has not been computed yet.
        pipe.reset()
        sensor.release_flow_gate.set()
        assert _wait_until(lambda: sensor._flow_calls >= 3), (
            "flow stage did not compute the first post-reset frame"
        )
        sensor.release_hhd_gate.set()
        deadline = time.perf_counter() + 5.0
        while len(frames) < count and time.perf_counter() < deadline:
            time.sleep(0.001)
    finally:
        sensor.release_flow_gate.set()
        sensor.release_hhd_gate.set()
        pipe.stop()
    assert len(frames) == count, f"expected {count} published frames, got {len(frames)}"


def test_frames_keep_the_reference_version_they_were_measured_with():
    """A reset queued after a frame's flow must not relabel that frame.

    Frames 0 and 1 are measured while the HHD stage is held on the first
    snapshot. The caller then queues a reset, and frame 2 is measured with the
    new version. Each snapshot must carry the version its own flow was computed
    with; reading the live counter at publish time would stamp the two older
    frames with the newer version.
    """
    sensor = SplitResetFakeSensor()
    frames: list[FrameSnapshot] = []
    _run_gated_reset(sensor, frames)
    versions = [snap.tracker_reset_seq for snap in frames]
    assert versions == [0, 0, 1], f"expected per-frame versions [0, 0, 1], got {versions}"


def test_reset_halves_run_on_the_threads_that_own_their_state():
    """Image state resets on the flow stage, force state on the HHD stage.

    The force half must land on the HHD stage before it decomposes the first
    frame of the new version. Resetting it from the flow thread - the old
    behaviour - mutates model/zero state that the HHD stage is using.
    """
    sensor = SplitResetFakeSensor()
    frames: list[FrameSnapshot] = []
    _run_gated_reset(sensor, frames)
    assert sensor.image_reset_threads == ["orisys-pipeline-flow"]
    assert sensor.force_reset_threads == ["orisys-pipeline-hhd"]
    # Two decompositions (frames 0 and 1) had completed when the force reset ran,
    # so it preceded the third one - the first frame of the new version.
    assert sensor.force_reset_at_hhd_calls == [2]
    assert sensor.hhd_calls == 3
    assert [snap.tracker_reset_seq for snap in frames] == [0, 0, 1]


def test_legacy_sensor_with_only_reset_is_reset_on_the_flow_stage():
    """A sensor without the split methods keeps the historical fallback."""
    sensor = LegacyResetFakeSensor()
    ready = threading.Event()
    release = threading.Event()
    state = {"produced": 0}

    def gated_produce():
        ready.set()
        release.wait(timeout=5.0)
        if state["produced"]:
            return None
        state["produced"] = 1
        return sensor.make_frame(0)

    frames: list = []
    pipe = StreamPipeline(sensor, gated_produce, on_frame=frames.append)
    pipe.start()
    try:
        assert ready.wait(timeout=5.0)
        pipe.reset()
        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)
    finally:
        release.set()
        pipe.stop()
    assert len(frames) == 1
    assert sensor.reset_calls == 1
    assert sensor.reset_threads == ["orisys-pipeline-flow"]
    assert sensor.reset_at_frame == [0], "the fallback reset must precede compute_flow"