orisys-app/shared/tests/test_history_source.py

"""Long-history browsing: interval reads, envelopes and one curve interface.

The replay used to load every row of a recording before drawing anything, so a
multi-hour session could not be shown at all. These tests pin the two things that
make it possible — an interval read bounded by the interval, and a whole-session
envelope built while streaming the file once — and that the charts can read from
either a session store or a recording through the same interface.
"""

from __future__ import annotations

import shutil
from pathlib import Path

import numpy as np
import pytest

from orisys_app_shared.acquisition.sdk_worker import SdkPacket
from orisys_app_shared.force.session import SessionSample, SessionStore
from orisys_app_shared.force.source import RecordingCurveSource, StoreCurveSource
from orisys_app_shared.record import ExperimentRecorder, RecordConfig, RecordingReader
from orisys_app_shared.record.index import RecordingIndex

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

ROWS = 12_000
SPIKE_ROW = 9_000
SPIKE_VALUE = 999.0


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


def _packet(index: int) -> SdkPacket:
    frame = np.zeros((4, 4, 3), np.uint8)
    fx = SPIKE_VALUE if index == SPIKE_ROW else 0.5
    return SdkPacket(
        raw_frame=frame,
        processed_frame=frame,
        timestamp=1000.0 + index / 30.0,
        mono_t=index / 30.0,
        fps=30.0,
        process_ms=1.0,
        normal_force=1.0,
        shear_x=fx,
        shear_y=-0.5,
        normal_force_raw=10.0 + (index % 7) * 0.01,
        shear_x_raw=fx,
        shear_y_raw=-0.5,
        force_unit="N",
        force_calibrated=True,
        frame_id=index,
    )


def _recording() -> Path:
    """A recording of ``ROWS`` rows with no drops: the writer is allowed to keep up."""
    import time

    root = _root()
    recorder = ExperimentRecorder(RecordConfig(video=False))
    directory = recorder.start(root=root, session_id=1, metadata={})
    for index in range(ROWS):
        recorder.on_packet(_packet(index))
        if index % 256 == 0:
            # Feeding 12000 rows in a tight loop would (correctly) overflow the
            # ordered stream and be recorded as gaps; this test is about history
            # browsing, so let the writer drain.
            deadline = time.monotonic() + 5.0
            while recorder._queue.qsize() > 512 and time.monotonic() < deadline:
                time.sleep(0.001)
    report = recorder.stop()
    assert report.stats.dropped_data == 0, report.problems
    assert report.stats.frames == ROWS
    return Path(directory)


def test_the_index_reads_only_the_interval_it_is_asked_for():
    directory = _recording()
    index = RecordingIndex(directory)
    index.build()
    assert index.rows == ROWS
    assert index.duration_s() == pytest.approx((ROWS - 1) / 30.0, abs=1e-6)

    # The whole range keeps every row, including the very first one.
    whole = index.read_rows(0, ROWS)
    assert len(whole) == ROWS
    assert whole.frame_ids[0] == 0
    assert whole.frame_ids[-1] == ROWS - 1

    start, stop = index.row_range(300.0, 300.5)
    # One row before the interval is included on purpose: a cursor exactly on a
    # sample still wants the row it follows.
    assert 15 <= stop - start <= 20, (start, stop)
    window = index.read_rows(start, stop)
    assert 15 <= len(window) <= 20
    assert window.times[0] >= 300.0 - 1 / 30.0
    assert window.times[-1] <= 300.5


def test_a_large_recording_is_loaded_as_a_marked_decimated_view():
    directory = _recording()
    reader = RecordingReader(directory, max_rows=500)
    info = reader.info()
    assert info.decimated, "a 12000-row recording was loaded in full"
    assert info.total_rows == ROWS
    assert len(info.times) <= 500
    # The duration comes from the frame index, never from the decimated view.
    assert info.duration_s == pytest.approx((ROWS - 1) / 30.0, abs=0.05)
    assert len(info.frames) == ROWS


def test_a_small_recording_is_not_decimated():
    directory = _recording()
    reader = RecordingReader(directory, max_rows=50_000)
    info = reader.info()
    assert not info.decimated
    assert len(info.times) == ROWS
    assert info.total_rows == ROWS


def test_an_interval_read_sees_every_row_including_a_single_peak():
    directory = _recording()
    reader = RecordingReader(directory, max_rows=500)
    spike_t = SPIKE_ROW / 30.0
    window = reader.window(spike_t - 0.05, spike_t + 0.05)
    assert window.times.size >= 3
    assert float(np.nanmax(window.value_for_unit("N")[:, 1])) == pytest.approx(SPIKE_VALUE)
    # ...and it is the *exact* interval, not the decimated whole-session view.
    assert not window.decimated
    assert window.times[0] >= spike_t - 0.1


def test_a_decimated_window_still_keeps_the_peak():
    directory = _recording()
    reader = RecordingReader(directory, max_rows=500)
    window = reader.window(0.0, (ROWS - 1) / 30.0, max_points=300, unit="N")
    assert window.decimated
    assert len(window) <= 600
    assert float(np.nanmax(window.value_for_unit("N")[:, 1])) == pytest.approx(SPIKE_VALUE)


def test_the_envelope_covers_the_whole_session_in_few_points():
    directory = _recording()
    reader = RecordingReader(directory, max_rows=500)
    overview = reader.overview(unit="N", buckets=200)
    assert 0 < len(overview) <= 200
    assert not np.isnan(overview.highs[:, 1]).all()
    # The spike is inside a bucket the strip would otherwise skip.
    assert float(np.nanmax(overview.highs[:, 1])) == pytest.approx(SPIKE_VALUE)
    assert float(overview.counts.sum()) > 0
    # Cached: asking again does not re-stream the file.
    assert reader.overview(unit="N", buckets=200) is overview


def test_the_two_sources_answer_the_same_questions():
    directory = _recording()
    reader = RecordingReader(directory, max_rows=ROWS + 1)
    recording = RecordingCurveSource(reader)

    store = SessionStore()
    for index in range(ROWS):
        sample = SessionSample(
            index,
            index / 30.0,
            fn_raw=10.0 + (index % 7) * 0.01,
            fx_raw=SPIKE_VALUE if index == SPIKE_ROW else 0.5,
            fy_raw=-0.5,
            fn_n=1.0,
            fx_n=SPIKE_VALUE if index == SPIKE_ROW else 0.5,
            fy_n=-0.5,
        )
        store.append(sample)
    memory = StoreCurveSource(store)

    for source in (recording, memory):
        first, last = source.extent()
        assert first == pytest.approx(0.0, abs=1e-6)
        assert last == pytest.approx((ROWS - 1) / 30.0, abs=1e-6)
        assert source.duration_s() == pytest.approx((ROWS - 1) / 30.0, abs=1e-6)
    assert recording.kind == "disk" and memory.kind == "memory"

    # The same interval, read from disk and from memory, agrees on the values.
    # (The disk read also includes the row just before the interval.)
    disk = recording.window(200.0, 210.0, unit="N")
    cached = memory.window(200.0, 210.0, unit="N")
    inside = (disk.times >= 200.0) & (disk.times <= 210.0)
    assert inside.sum() == len(cached) > 0
    assert np.allclose(disk.values[inside], cached.values, equal_nan=True)
    # The CSV keeps six decimals; the in-memory store keeps the raw float.
    assert np.allclose(disk.times[inside], cached.times, atol=1e-6)

    # ...and so do the statistics of that interval (the disk read may include the
    # one row that precedes it).
    a = recording.statistics(200.0, 210.0, unit="N")
    b = memory.statistics(200.0, 210.0, unit="N")
    assert abs(a.samples - b.samples) <= 1
    assert a.peaks[1] == pytest.approx(b.peaks[1])
    assert a.means[1] == pytest.approx(b.means[1], abs=1e-6)


def test_the_whole_session_can_be_read_without_holding_it_all():
    """The point of the index: peak value with a bounded number of rows held."""
    directory = _recording()
    reader = RecordingReader(directory, max_rows=500)
    overview = reader.overview(unit="N", buckets=100)
    # What a whole-session strip needs is one band per bucket, not 12000 rows.
    assert len(overview) <= 100
    assert len(reader.info().times) <= 500
    assert reader.info().total_rows == ROWS
    # A single sample can still be read exactly, on demand.
    spike_t = SPIKE_ROW / 30.0
    exact = reader.window(spike_t, spike_t + 1 / 30.0)
    assert 1 <= len(exact) <= 3, len(exact)
    assert float(np.nanmax(exact.newtons[:, 1])) == pytest.approx(SPIKE_VALUE)