tests/test_app_helpers.py

"""Unit tests for streaming helpers and the Gradio fork link. No GPU or weights."""

from __future__ import annotations

import os
import pickle
import threading
import traceback
import uuid
from collections.abc import Callable
from fractions import Fraction
from pathlib import Path

import numpy as np
import pytest
import rerun as rr
import rerun.blueprint as rrb
import torch

os.environ.setdefault("FDANYONE_SKIP_LOAD", "1")

from fdanyone.config import INFERENCE  # noqa: E402
from fdanyone.errors import FourDAnyoneError  # noqa: E402
from fdanyone.model.inference import _tensor_frames  # noqa: E402

import fdanyone.rerun_streaming as streaming  # noqa: E402
import fdanyone.viz as viz  # noqa: E402
import fdanyone_app  # noqa: E402

EXAMPLE: Path = Path(__file__).resolve().parents[1] / "examples" / "jump-rope.mp4"


def test_preview_slice_plan_covers_every_fourth_latent_frame() -> None:
    plan: tuple[tuple[int, int], ...] = streaming.preview_slice_plan(31)
    assert plan == (
        (0, 0),
        (4, 16),
        (8, 32),
        (12, 48),
        (16, 64),
        (20, 80),
        (24, 96),
        (28, 112),
    )
    # Eight previews per view, and none past the frozen 121-frame contract.
    assert len(plan) == 8
    assert max(source for _, source in plan) < INFERENCE.num_frames


def test_preview_slice_plan_rejects_empty_latents() -> None:
    with pytest.raises(ValueError):
        streaming.preview_slice_plan(0)


def test_tensor_frames_maps_the_signed_range_onto_uint8() -> None:
    """The preview reuses the pipeline's own truncation, so pin its endpoints."""

    video: torch.Tensor = torch.tensor([-1.0, 0.0, 1.0]).reshape(3, 1, 1, 1).expand(3, 1, 1, 1)
    frame: np.ndarray = next(iter(_tensor_frames(video)))
    assert frame.dtype == np.uint8
    assert frame.shape == (1, 1, 3)
    assert frame.reshape(3).tolist() == [0, 127, 255]


def test_probe_clip_accepts_the_bundled_example() -> None:
    info: streaming.ClipInfo = streaming.probe_clip(EXAMPLE, 0.0)
    assert info.fps > 0
    assert info.duration_seconds >= info.required_seconds
    assert info.required_seconds == pytest.approx(float(Fraction(120, 1) / info.fps))


def test_probe_clip_rejects_a_start_time_past_the_end() -> None:
    with pytest.raises(FourDAnyoneError, match="need"):
        streaming.probe_clip(EXAMPLE, 3600.0)


def test_probe_clip_rejects_a_missing_file(tmp_path: Path) -> None:
    with pytest.raises(FourDAnyoneError, match="does not exist"):
        streaming.probe_clip(tmp_path / "absent.mp4", 0.0)


def _pump_session(data_dir: Path = Path("/nonexistent")) -> streaming.Session:
    """A Session with only the fields ``pump`` touches."""

    return streaming.Session(
        streaming.RunSpec(
            token="0" * 32,
            video_path=EXAMPLE,
            start_time=0.0,
            seed=0,
            fps=Fraction(25, 1),
            data_dir=data_dir,
            created=0.0,
        )
    )


def test_pump_yields_hook_labels_and_returns_the_work_result() -> None:
    session: streaming.Session = _pump_session()

    def work() -> str:
        for label in ("bboxes", "smplx"):
            session.events.put(label)
        return "prepared"

    pump = streaming.pump(session, work)
    seen: list[str] = []
    while True:
        try:
            seen.append(next(pump))
        except StopIteration as done:
            assert done.value == "prepared"
            break
    assert seen == ["bboxes", "smplx"]


def test_stop_sentinel_unwinds_a_running_pump_worker(tmp_path: Path) -> None:
    """Stop must reach the pipeline thread itself, not just Gradio's generator.

    ``cancels`` closes the request generator, but the worker thread (or the
    forked ZeroGPU child) never hears it; the sentinel checked by the hooks is
    the only thing that actually frees the GPU. The barrier makes the order
    deterministic: the worker is provably mid-run when Stop arrives.
    """

    assert streaming.request_stop(None) == "Nothing is running."

    session: streaming.Session = _pump_session(tmp_path)
    stop_requested: threading.Event = threading.Event()

    def work() -> str:
        # A pipeline hook checks the sentinel once per stage; model the stage
        # boundary the run is inside when the user presses Stop.
        streaming.check_stop(session.spec)
        session.events.put("running")
        assert stop_requested.wait(timeout=5.0)
        streaming.check_stop(session.spec)
        return "finished"

    pump = streaming.pump(session, work)
    assert next(pump) == "running"
    streaming.request_stop(session.spec)
    assert (tmp_path / "stop-requested").exists()
    stop_requested.set()
    with pytest.raises(streaming.RunCancelled):
        list(pump)


def test_pump_reraises_a_worker_failure_on_the_caller_thread() -> None:
    """A pipeline error must surface where Gradio can turn it into a message."""

    session: streaming.Session = _pump_session()

    def work() -> str:
        raise FourDAnyoneError("boom")

    with pytest.raises(FourDAnyoneError, match="boom"):
        list(streaming.pump(session, work))


def test_blueprints_build_for_every_phase() -> None:
    """Blueprints expose every phase in the approved camera-centered layout."""

    assert streaming.APPLICATION_ID == "4danyone-rerun-v3"
    assert viz.APPLICATION_ID == "4danyone-rerun-v3"

    motion: rrb.Blueprint = viz.motion_blueprint()
    source: rrb.Spatial2DView = motion.root_container.contents[0]  # pyrefly: ignore
    detections: rrb.Spatial2DView = motion.root_container.contents[1]  # pyrefly: ignore
    assert (source.origin, source.contents) == ("world/camera", ["$origin/video"])
    assert detections.origin == "world/camera"

    assert viz.diffusion_blueprint() is not None

    body: rrb.Blueprint = viz.body_blueprint()
    body_world: rrb.Spatial3DView = body.root_container.contents[0]  # pyrefly: ignore
    body_source: rrb.Spatial2DView = body.root_container.contents[1]  # pyrefly: ignore
    assert (body_world.origin, body_world.name) == ("world", "Canonical world")
    assert body_source.contents == ["$origin/**", "/world/body/**"]
    assert body.time_panel.play_state == "playing"

    result: rrb.Blueprint = viz.result_blueprint(Fraction(25, 1))
    top: rrb.Horizontal = result.root_container.contents[0]  # pyrefly: ignore
    dense: rrb.Horizontal = result.root_container.contents[1]  # pyrefly: ignore
    world: rrb.Spatial3DView = top.contents[0]  # pyrefly: ignore
    result_source: rrb.Spatial2DView = top.contents[1]  # pyrefly: ignore
    assert result.root_container.row_shares == [2.0, 1.0]
    assert top.column_shares == [3.0, 1.0]
    assert (world.origin, world.name) == ("world", "Canonical world")
    assert (result_source.origin, result_source.contents) == (
        "world/camera",
        ["$origin/**", "/world/body/**"],
    )
    assert type(dense).__name__ == "Horizontal"
    assert tuple(view.origin for view in dense.contents) == tuple(
        f"world/cameras/dense/{index:02d}/image" for index in range(viz.VIEWS)
    )
    assert result.time_panel.play_state == "playing"
    assert result.time_panel.loop_mode == "all"


# ---------------------------------------------------------------------------
# Fork safety — the one thing that only misbehaves on ZeroGPU
# ---------------------------------------------------------------------------

OK: int = 0
"""Child exit code meaning the behaviour under test held."""

WRONG: int = 3
"""Child exit code meaning it did not; distinct from any interpreter failure."""


def _in_child(body: Callable[[], int]) -> int:
    """Run ``body`` in a forked child and return its exit code.

    ``os._exit`` skips every atexit handler and any pytest teardown the child
    inherited, so only what ``body`` decides reaches the parent.
    """

    pid: int = os.fork()
    if pid == 0:
        try:
            code: int = body()
        except BaseException:  # noqa: BLE001 - a traceback here is the diagnosis
            traceback.print_exc()
            code = WRONG
        os._exit(code)
    _, status = os.waitpid(pid, 0)
    return os.waitstatus_to_exitcode(status)


def test_a_recording_made_before_a_fork_cannot_be_flushed_after_it() -> None:
    """Pin the ZeroGPU failure the per-link recording exists to avoid.

    ``@spaces.GPU`` runs its callback in a forked child, so a recording opened
    by an earlier link is a parent's recording. The SDK notices the pid change
    and refuses: "Fork detected during flush". Everything ran fine locally
    before this was understood, because ``spaces`` never forks off a Space.
    """

    recording = rr.RecordingStream(streaming.APPLICATION_ID, recording_id=uuid.uuid4().hex)
    stream = recording.binary_stream()
    recording.log("log", rr.TextLog("parent"))
    assert stream.read()

    def body() -> int:
        recording.log("log", rr.TextLog("child"))
        try:
            stream.read()
        except RuntimeError as exc:
            return OK if "fork" in str(exc).lower() else WRONG
        return WRONG

    assert _in_child(body) == OK, (
        "A pre-fork recording flushed cleanly from the child. If the Rerun SDK "
        "now supports this, the per-link recording is merely unnecessary, not wrong."
    )


def test_open_link_gives_a_forked_child_a_recording_it_can_flush(tmp_path: Path) -> None:
    """The fix: each link builds its own stream, in whatever process it runs in.

    The token is the ``recording_id``, so the parent's rows and the child's are
    one recording to the viewer even though no object crossed the fork.
    """

    token: str = uuid.uuid4().hex
    parent: fdanyone_app.Link = fdanyone_app.open_link(token, tmp_path / "parent.rrd")
    parent.recording.log("log", rr.TextLog("parent"))
    assert parent.read()
    parent.close()

    def body() -> int:
        child: fdanyone_app.Link = fdanyone_app.open_link(token, tmp_path / "child.rrd")
        assert child.recording.get_recording_id() == token
        child.recording.log("log", rr.TextLog("child"))
        payload: bytes | None = child.read()
        return OK if payload else WRONG

    assert _in_child(body) == OK
    # The dual sink teed the parent's rows into its part file as well.
    assert (tmp_path / "parent.rrd").stat().st_size > 0


RRD_MAGIC: bytes = b"RRF2"
"""First four bytes of every RRD document the SDK emits."""


def test_a_stream_per_link_adds_no_framing_the_viewer_did_not_already_get(tmp_path: Path) -> None:
    """Why splitting one stream into four costs the browser viewer nothing.

    Each ``read`` already returns a whole RRD document, magic bytes and manifest
    included, so what the app sends the viewer has always been a concatenation
    of documents rather than one long one. Opening a stream per link only
    changes how many of them carry the same store id. If this ever stops
    holding — if a later read became a bare continuation of the first — then
    per-link streams would start sending headers mid-stream, and this test is
    where that shows up rather than in a blank viewer on the Space.
    """

    link: fdanyone_app.Link = fdanyone_app.open_link(uuid.uuid4().hex, tmp_path / "part.rrd")
    for row in range(3):
        link.recording.log("log", rr.TextLog(f"row {row}"))
        payload: bytes | None = link.read()
        assert payload is not None
        assert payload.startswith(RRD_MAGIC), f"read {row} is not a self-contained document"


def test_coco17_annotation_context_matches_the_vitpose_layout() -> None:
    """The context carries simplecv's COCO-17 body joints and their links."""

    from simplecv.data.skeleton.coco_133 import COCO_17_IDX, COCO_133_ID2NAME

    assert len(COCO_17_IDX) == 17
    assert COCO_133_ID2NAME[0] == "nose"
    assert COCO_133_ID2NAME[16] == "right_ankle"
    context: rr.AnnotationContext = viz.coco17_annotation_context()
    assert context is not None


def test_format_status_counts_known_stages_and_passes_failures_through() -> None:
    """The banner shows x/x plus elapsed for the fixed run shape, nothing else."""

    spec: streaming.RunSpec = streaming.new_spec(EXAMPLE, 0.0, 0)
    first: str = streaming.format_status(spec, streaming.STAGES[0])
    assert first.startswith(f"{streaming.STAGES[0]} (1/{len(streaming.STAGES)}")
    last: str = streaming.format_status(spec, streaming.STAGES[-1])
    assert f"({len(streaming.STAGES)}/{len(streaming.STAGES)}" in last
    assert streaming.format_status(spec, "Failed: boom") == "Failed: boom"


def test_stage_sequence_is_complete_and_unique() -> None:
    """One count per stage: every emitted label indexes, none twice.

    ``STAGES`` is assembled from the same registries the hooks emit from, so
    this pins the whole contract: 19 distinct stages, every registry label
    present, and the denoise labels formatted by the shared helper.
    """

    assert len(streaming.STAGES) == 19
    assert len(set(streaming.STAGES)) == len(streaming.STAGES)
    for label in streaming.MOTION_LABELS.values():
        assert label in streaming.STAGE_INDEX
    for label in streaming.GENERATE_LABELS.values():
        assert label in streaming.STAGE_INDEX
    for step in range(1, streaming.SETTINGS.num_inference_steps + 1):
        assert streaming.denoise_step_label(step) in streaming.STAGE_INDEX


def test_a_run_spec_survives_the_pickle_into_a_zerogpu_worker() -> None:
    """Arguments are all a GPU callback gets; unpicklable state never arrives."""

    spec: streaming.RunSpec = streaming.new_spec(EXAMPLE, 0.0, 7)
    restored: streaming.RunSpec = pickle.loads(pickle.dumps(spec))
    assert restored == spec
    assert restored.data_dir.name == spec.token