video-image-production-build/tests/observability/test_pipeline_context.py

from __future__ import annotations

from pathlib import Path
from types import SimpleNamespace

from obagent_sdk import observe

from production_build_agents.observability.pipeline import (
    PipelineObservation,
)


class _Handle:
    def __init__(
        self,
        *,
        run_id: str | None = None,
        inst_id: str | None = None,
    ) -> None:
        self.run_id = run_id
        self.inst_id = inst_id
        self.declared: list[dict] = []
        self.outputs: list[tuple[object, dict]] = []
        self.finished: list[dict] = []

    def declare(self, **payload) -> None:
        self.declared.append(payload)

    def set_output(self, output, **payload) -> None:
        self.outputs.append((output, payload))

    def finish(self, **payload) -> None:
        self.finished.append(payload)


class _Manager:
    def __init__(self, handle: _Handle) -> None:
        self.handle = handle
        self.exits: list[tuple] = []

    def __enter__(self) -> _Handle:
        return self.handle

    def __exit__(self, *args):
        self.exits.append(args)
        if args[1] is not None:
            raise args[1]
        return False


def test_pipeline_has_two_stages_and_exposes_child_parent_ids(
    monkeypatch,
    tmp_path: Path,
) -> None:
    source = tmp_path / "production_final.json"
    anchor = tmp_path / "anchor.png"
    source.write_text('{"title":"真实输入"}', encoding="utf-8")
    anchor.write_bytes(b"anchor")
    run_handle = _Handle(run_id="pipeline-run-id")
    run_manager = _Manager(run_handle)
    stage_handles = iter(
        (
            _Handle(inst_id="global-stage-inst"),
            _Handle(inst_id="production-stage-inst"),
        )
    )
    stage_managers: list[_Manager] = []
    captured: dict = {}

    def fake_run(**kwargs):
        captured.update(kwargs)
        return run_manager

    def fake_module(*_args, **_kwargs):
        manager = _Manager(next(stage_handles))
        stage_managers.append(manager)
        return manager

    monkeypatch.setattr(observe, "run", fake_run)
    monkeypatch.setattr(observe, "module", fake_module)
    settings = SimpleNamespace(
        configure_sdk=lambda: True,
        project="pipeline-project",
        project_name="Pipeline 观测",
        capture_mode="full",
    )

    with PipelineObservation(
        project_root=tmp_path,
        round_id="round-1",
        input_path=source,
        output_root=tmp_path / "runs",
        shared_visual_anchor_path=anchor,
        settings=settings,
    ) as pipeline:
        with pipeline.stage(
            "global_data",
            input_value={"thread_id": "round-1-global-data"},
        ) as stage:
            assert stage.parent_run_id == "pipeline-run-id"
            assert stage.parent_inst_id == "global-stage-inst"
            stage.finish(
                {
                    "run_id": "round-1-global-data",
                    "status": "COMPLETED",
                    "global_data_delivery_path": "/run/delivery.json",
                }
            )
        with pipeline.stage(
            "production",
            input_value={"thread_id": "round-1-production"},
        ) as stage:
            assert stage.parent_run_id == "pipeline-run-id"
            assert stage.parent_inst_id == "production-stage-inst"
            stage.finish(
                {
                    "run_id": "round-1-production",
                    "status": "RUNNING",
                }
            )
        pipeline.finish(
            {
                "round_id": "round-1",
                "status": "RUNNING",
                "phase": "PRODUCTION",
                "global_data_thread_id": "round-1-global-data",
                "production_thread_id": "round-1-production",
            }
        )

    assert captured["agent"] == "pipeline"
    assert captured["auto_collect"] is False
    assert [node["key"] for node in captured["spec"]["nodes"]] == [
        "global_data",
        "production",
    ]
    assert captured["spec"]["edges"] == [
        {"source": "global_data", "target": "production"}
    ]
    global_output = stage_managers[0].handle.outputs[0][0]
    assert global_output["阶段"] == "Global Data"
    assert global_output["执行结论"] == "已完成"
    assert global_output["正式交付"] == "/run/delivery.json"
    global_input = stage_managers[0].handle.declared[0]["blocks"][0].value
    assert global_input["阶段任务"] == (
        "整理 Production 所需的全局资料与约束"
    )
    assert "thread_id" not in global_input
    final_output = run_handle.finished[0]["final_output"]
    assert final_output["当前阶段"] == "Production"
    assert final_output["执行结论"] == "进行中"
    assert "phase" not in final_output