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

from __future__ import annotations

from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from io import BytesIO
import tempfile
import threading
import unittest
from pathlib import Path
from urllib.request import urlopen

from PIL import Image

from production_build_agents.contracts.models import (
    Artifact,
    ToolArtifactMapping,
    ToolCallRecord,
)
from production_build_agents.run.artifacts import (
    ArtifactIntegrityError,
    content_identity_for_bytes,
    content_identity_for_path,
    resolve_candidate_artifact_path,
    seal_shared_visual_anchor,
    verify_artifact,
    verify_shared_visual_anchor,
)
from production_build_agents.tools.media import probe_media


def _png_bytes(color: str) -> bytes:
    output = BytesIO()
    Image.new("RGB", (2, 2), color=color).save(output, format="PNG")
    return output.getvalue()


class _MutableImageHandler(BaseHTTPRequestHandler):
    body = _png_bytes("red")

    def do_GET(self) -> None:
        self.send_response(200)
        self.send_header("Content-Type", "image/png")
        self.send_header("Content-Length", str(len(self.body)))
        self.end_headers()
        self.wfile.write(self.body)

    def log_message(self, format: str, *args: object) -> None:
        del format, args


class ArtifactIntegrityTest(unittest.TestCase):
    def test_ass_uses_stable_document_mime(self) -> None:
        self.assertEqual(
            content_identity_for_bytes(
                b"[Script Info]\n",
                filename="subtitles.ass",
            )["mime_type"],
            "text/x-ass",
        )
        with tempfile.TemporaryDirectory() as temp_dir:
            path = Path(temp_dir) / "subtitles.ass"
            path.write_text("[Script Info]\n", encoding="utf-8")
            self.assertEqual(
                content_identity_for_path(path)["mime_type"],
                "text/x-ass",
            )

    def test_remote_candidate_resolves_to_single_local_copy(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            local = Path(temp_dir) / "sealed.png"
            local.write_bytes(b"sealed")
            remote = "https://example.test/mutable.png"
            record = ToolCallRecord(
                tool_call_id="view-1",
                tool_name="view_images",
                success=True,
                artifact_mappings=[
                    ToolArtifactMapping(
                        source=remote,
                        local_path=str(local),
                    )
                ],
            )
            self.assertEqual(
                resolve_candidate_artifact_path(remote, [record]),
                local.resolve(),
            )

    def test_formal_remote_uri_is_rejected(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            cached = Path(temp_dir) / "old.png"
            cached.write_bytes(b"old-sealed-content")
            remote = "https://example.test/mutable.png"
            artifact = Artifact(
                artifact_id="Task1-v1-artifact-1",
                artifact_type="image",
                uri=remote,
                source_uri=remote,
                description="不得作为正式远程 URI",
                **content_identity_for_path(cached),
            )
            with self.assertRaisesRegex(
                ArtifactIntegrityError,
                "必须是本地封印路径",
            ):
                verify_artifact(artifact)

    def test_shared_visual_anchor_rejects_local_content_replacement(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            path = Path(temp_dir) / "anchor.png"
            path.write_bytes(_png_bytes("red"))
            anchor = seal_shared_visual_anchor(path)

            path.write_bytes(_png_bytes("blue"))

            with self.assertRaisesRegex(
                ArtifactIntegrityError,
                "视觉锚点内容已变化",
            ):
                verify_shared_visual_anchor(anchor)

    def test_localized_artifact_survives_remote_url_content_replacement(
        self,
    ) -> None:
        server = ThreadingHTTPServer(("127.0.0.1", 0), _MutableImageHandler)
        thread = threading.Thread(target=server.serve_forever, daemon=True)
        thread.start()
        try:
            with tempfile.TemporaryDirectory() as temp_dir:
                remote = (
                    f"http://127.0.0.1:{server.server_port}/mutable.png"
                )
                original_remote = _png_bytes("red")
                replacement_remote = _png_bytes("blue")
                _MutableImageHandler.body = original_remote
                result = probe_media(
                    remote,
                    output_dir=Path(temp_dir),
                )
                localized = Path(result["local_path"])
                artifact = Artifact(
                    artifact_id="Task1-v1-artifact-1",
                    artifact_type="image",
                    uri=str(localized),
                    source_uri=remote,
                    description="正式产物使用首次本地化并封印的内容",
                    **content_identity_for_path(localized),
                )

                _MutableImageHandler.body = replacement_remote
                with urlopen(remote, timeout=2) as response:
                    self.assertEqual(response.read(), replacement_remote)

                self.assertEqual(localized.read_bytes(), original_remote)
                verify_artifact(artifact)
        finally:
            server.shutdown()
            server.server_close()
            thread.join(timeout=2)


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