video-image-production-build/tests/run/test_segment_inputs.py

from __future__ import annotations

import tempfile
import unittest
from pathlib import Path

from production_build_agents.contracts.document_identity import (
    exact_document_sha256,
)
from production_build_agents.contracts.production_execution_models import (
    ImmutableJsonRef,
)
from production_build_agents.contracts.production_planning_models import (
    ProductionPlanningPackage,
)
from production_build_agents.run.artifacts import (
    content_identity_for_path,
    seal_shared_visual_anchor,
)
from production_build_agents.run.segment_inputs import (
    SegmentInputIntegrityError,
    accepted_segment_snapshot,
    load_verified_segment_inputs,
)
from tests.support.production_fixtures import (
    production_plan,
    segment_bundle,
)
from tests.support.production_planner_fixtures import (
    build_production_planner_scenario,
)


def _ref(path: Path) -> ImmutableJsonRef:
    return ImmutableJsonRef(
        uri=str(path.resolve()),
        document_sha256=exact_document_sha256(path.read_bytes()),
    )


class SegmentInputsTest(unittest.TestCase):
    def _package_with_global_pool(self, root: Path):
        package = segment_bundle(1)[0]
        artifacts = []
        for index, base in enumerate(
            (
                package.artifacts[0],
                package.artifacts[0].model_copy(
                    update={
                        "artifact_id": "Task2-v1-artifact-1",
                        "artifact_type": "audio",
                        "description": "unplanned optional audio",
                    }
                ),
            ),
            start=1,
        ):
            path = root / f"global-{index}.bin"
            path.write_bytes(f"global-{index}".encode())
            artifacts.append(
                base.model_copy(
                    update={
                        "uri": str(path.resolve()),
                        "source_uri": f"https://mutable.test/{index}",
                        **content_identity_for_path(path),
                    }
                )
            )
        delivery = build_production_planner_scenario(
            run_id="segment-input-global"
        ).global_data_delivery.model_copy(
            update={"active_artifacts": artifacts}
        )
        delivery_path = root / "global_data_stage_delivery.json"
        delivery_path.write_text(delivery.model_dump_json(), encoding="utf-8")
        anchor_path = root / "anchor.png"
        anchor_path.write_bytes(b"anchor")
        package = package.model_copy(
            update={
                "artifacts": [artifacts[0]],
                "global_data_delivery_ref": _ref(delivery_path),
                "shared_visual_anchor": seal_shared_visual_anchor(anchor_path),
            }
        )
        return package, delivery_path, artifacts

    def test_all_active_artifacts_are_available_not_only_planner_selection(
        self,
    ) -> None:
        with tempfile.TemporaryDirectory() as directory:
            package, _delivery_path, artifacts = self._package_with_global_pool(
                Path(directory)
            )
            verified = load_verified_segment_inputs(package)
            self.assertEqual(
                [item.artifact_id for item in verified.global_artifacts],
                [item.artifact_id for item in artifacts],
            )
            self.assertEqual(
                [
                    item.artifact_id
                    for item in verified.lineage_artifacts(package)
                ],
                ["Task2-v1-artifact-1"],
            )
            payload = verified.agent_payload()
            self.assertNotIn("tool_calls", str(payload))
            self.assertNotIn("source_uri", str(payload))
            self.assertNotIn("https://mutable.test/", str(payload))

    def test_exact_delivery_and_selected_artifact_identity_fail_closed(
        self,
    ) -> None:
        with tempfile.TemporaryDirectory() as directory:
            package, delivery_path, artifacts = (
                self._package_with_global_pool(Path(directory))
            )
            delivery_path.write_text("{}", encoding="utf-8")
            with self.assertRaisesRegex(
                SegmentInputIntegrityError,
                "精确字节已变化",
            ):
                load_verified_segment_inputs(package)

    def test_global_pool_rejects_same_id_or_uri_with_different_identity(
        self,
    ) -> None:
        for collision in ("artifact_id", "uri"):
            with self.subTest(collision=collision):
                with tempfile.TemporaryDirectory() as directory:
                    root = Path(directory)
                    package, delivery_path, artifacts = (
                        self._package_with_global_pool(root)
                    )
                    delivery = build_production_planner_scenario(
                        run_id="collision-global"
                    ).global_data_delivery.model_copy(
                        update={
                            "active_artifacts": [
                                artifacts[0],
                                artifacts[1].model_copy(
                                    update={
                                        collision: getattr(
                                            artifacts[0],
                                            collision,
                                        )
                                    }
                                ),
                            ]
                        }
                    )
                    delivery_path.write_text(
                        delivery.model_dump_json(),
                        encoding="utf-8",
                    )
                    package = package.model_copy(
                        update={
                            "global_data_delivery_ref": _ref(delivery_path)
                        }
                    )
                    with self.assertRaisesRegex(
                        SegmentInputIntegrityError,
                        "必须唯一|存在歧义",
                    ):
                        load_verified_segment_inputs(package)

    def test_global_pool_rejects_symlink_uri_alias(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            package, delivery_path, artifacts = (
                self._package_with_global_pool(root)
            )
            alias = root / "same-artifact-alias.bin"
            alias.symlink_to(Path(artifacts[0].uri))
            delivery = build_production_planner_scenario(
                run_id="alias-global"
            ).global_data_delivery.model_copy(
                update={
                    "active_artifacts": [
                        artifacts[0],
                        artifacts[1].model_copy(
                            update={
                                "uri": str(alias),
                                **content_identity_for_path(alias),
                            }
                        ),
                    ]
                }
            )
            delivery_path.write_text(
                delivery.model_dump_json(),
                encoding="utf-8",
            )
            package = package.model_copy(
                update={"global_data_delivery_ref": _ref(delivery_path)}
            )

            with self.assertRaisesRegex(
                SegmentInputIntegrityError,
                "URI.*歧义|uri 必须唯一",
            ):
                load_verified_segment_inputs(package)

    def test_selected_artifact_must_equal_global_delivery_record(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            package, _delivery_path, artifacts = (
                self._package_with_global_pool(Path(directory))
            )
            package = package.model_copy(
                update={
                    "artifacts": [
                        artifacts[0].model_copy(
                            update={"description": "forged"}
                        )
                    ]
                }
            )
            with self.assertRaisesRegex(
                SegmentInputIntegrityError,
                "不等于正式 Global Data",
            ):
                load_verified_segment_inputs(package)

    def test_snapshot_unions_planning_baseline_and_current_pass(self) -> None:
        first = segment_bundle(1)[-1]
        second = segment_bundle(2)[-1]
        packages = [segment_bundle(1)[0], segment_bundle(2)[0]]
        plan = production_plan(packages).model_copy(
            update={"regenerate_segment_ids": ["Segment2"]}
        )
        planning = ProductionPlanningPackage.model_construct(
            accepted_segments=[first, second]
        )
        refreshed_second = second.model_copy(
            update={"source_plan_version": 2}
        )
        snapshot = accepted_segment_snapshot(
            planning,
            plan,
            {"Segment2": refreshed_second},
            exclude_segment_id="Segment2",
        )
        self.assertEqual(list(snapshot), ["Segment1"])
        full = accepted_segment_snapshot(
            planning,
            plan,
            {"Segment2": refreshed_second},
        )
        self.assertIs(full["Segment2"], refreshed_second)


if __name__ == "__main__":
    unittest.main()