"""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