tests/test_pipeline_lifecycle.py

"""Scratch-lifecycle tests: ``generate_run`` consumes its ``PreparedRun``.

The cleanup guarantee used to live in every caller; these tests pin it to the
one owner. Settlement and discard must both happen exactly once whether
generation succeeds, generation raises (including a cancellation hook), or
settlement itself raises.
"""

from __future__ import annotations

import os
from fractions import Fraction
from pathlib import Path

import pytest

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

import fdanyone.model.inference
import fdanyone.output
from fdanyone.pipeline import PreparedRun, generate_run  # noqa: E402


class _Conditioning:
    """Stands in for target conditioning with a deferred skeleton render."""

    def __init__(self, settlement_error: Exception | None = None) -> None:
        self.target_waiter: object | None = object()
        self.settlement_error: Exception | None = settlement_error
        self.settled: int = 0

    def wait_for_target_skeletons(self) -> None:
        self.settled += 1
        if self.settlement_error is not None:
            raise self.settlement_error


def _prepared(tmp_path: Path, conditioning: _Conditioning) -> PreparedRun:
    scratch: Path = tmp_path / "scratch"
    scratch.mkdir()
    (scratch / "intermediate.bin").write_bytes(b"scratch")
    return PreparedRun(
        settings=object(),
        clip=object(),
        motion=object(),
        conditioning=conditioning,
        checkpoint=tmp_path / "checkpoint.safetensors",
        base_assets=object(),
        prompt_embedding_path=None,
        model_identity={},
        device="cpu",
        scratch=scratch,
        motion_dir=tmp_path / "motion",
        result_dir=tmp_path / "result",
        pipeline_started=0.0,
    )


def test_generate_run_settles_and_discards_on_success(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    conditioning: _Conditioning = _Conditioning()
    prepared: PreparedRun = _prepared(tmp_path, conditioning)
    monkeypatch.setattr(fdanyone.model.inference, "generate_views", lambda **_: object())
    monkeypatch.setattr(
        fdanyone.output,
        "export_result",
        lambda **_: {"fps": Fraction(25, 1), "total_pipeline_elapsed_seconds": 1.5},
    )

    result = generate_run(prepared, seed=0)

    assert result.result_dir == prepared.result_dir
    assert result.total_elapsed_seconds == 1.5
    assert result.metadata["result_dir"] == str(prepared.result_dir)
    assert conditioning.settled == 1
    assert not prepared.scratch.exists()


def test_generate_run_settles_and_discards_when_generation_raises(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    conditioning: _Conditioning = _Conditioning()
    prepared: PreparedRun = _prepared(tmp_path, conditioning)

    def explode(**_: object) -> object:
        raise RuntimeError("cancelled mid-denoise")

    monkeypatch.setattr(fdanyone.model.inference, "generate_views", explode)

    with pytest.raises(RuntimeError, match="cancelled mid-denoise"):
        generate_run(prepared, seed=0)
    assert conditioning.settled == 1
    assert not prepared.scratch.exists()


def test_generate_run_discards_even_when_settlement_raises(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    conditioning: _Conditioning = _Conditioning(settlement_error=RuntimeError("render worker died"))
    prepared: PreparedRun = _prepared(tmp_path, conditioning)
    monkeypatch.setattr(fdanyone.model.inference, "generate_views", lambda **_: object())
    monkeypatch.setattr(
        fdanyone.output,
        "export_result",
        lambda **_: {"total_pipeline_elapsed_seconds": 0.0},
    )

    with pytest.raises(RuntimeError, match="render worker died"):
        generate_run(prepared, seed=0)
    assert conditioning.settled == 1
    assert not prepared.scratch.exists()